diff --git a/benchmarks/README.md b/benchmarks/README.md index b5355f08..c9743ebc 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -123,8 +123,8 @@ Logs go to `logs/`. ## Environment - GPU: NVIDIA H100 SXM 80GB -- PyTorch 2.6.0+cu129 -- CUDA 12.9 +- PyTorch 2.10.0+cu130 +- CUDA 13.0 - Apex 0.1 (FusedLAMB) - QuACK 0.2.10 (fused RMSNorm) - NVIDIA DALI 1.53.0 diff --git a/benchmarks/benchmark_forward_time_nd_resolution.py b/benchmarks/benchmark_forward_time_nd_resolution.py new file mode 100644 index 00000000..535ce4e2 --- /dev/null +++ b/benchmarks/benchmark_forward_time_nd_resolution.py @@ -0,0 +1,810 @@ +# 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, + _gdp_mixer_cfg, + _hyena_mixer_cfg, + _mamba_mixer_cfg, +) + +from nvsubquadratic.lazy_config import instantiate + + +MIXER_CHOICES = ("attention", "flex", "fa4", "hyena", "mamba", "mamba_causal", "gdp") +# 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, + short_conv_backend: str = "subq_ops", + mamba_d_state: int | None = None, + mamba_ngroups: int | None = None, + gdp_householder: int = 3, + hyena_expansion: float = 1.0, +) -> 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), + short_conv_backend=short_conv_backend, + expansion=hyena_expansion, + ) + 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 in ("mamba", "mamba_causal"): + # Two series, not a flag: 'mamba' is bidirectional (two scans, the right + # vision/ND baseline) and 'mamba_causal' is unidirectional (one scan, what a + # language model runs). Bidirectional does ~2x the work, so which one a plot + # uses changes the Hyena comparison's sign — run both and let the reader pick. + cfg = _mamba_mixer_cfg( + hidden_dim, + headdim=mamba_headdim, + expand=mamba_expand, + bidirectional=(name == "mamba"), + d_state=mamba_d_state, + ngroups=mamba_ngroups, + ) + elif name == "gdp": + # GDP reads the same mamba_* knobs upstream, so it is sized from them here too: + # that is exactly what a drop-in --spec swap of a Mamba-2 layer produces. + cfg = _gdp_mixer_cfg( + hidden_dim, + num_householder=gdp_householder, + d_state=mamba_d_state if mamba_d_state is not None else 128, + headdim=mamba_headdim, + ngroups=mamba_ngroups if mamba_ngroups is not None else 8, + expand=mamba_expand, + ) + 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, + short_conv_backend: str, + mamba_d_state: int | None, + mamba_ngroups: int | None, + gdp_householder: int, + hyena_expansion: float, + 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, + short_conv_backend=short_conv_backend, + mamba_d_state=mamba_d_state, + mamba_ngroups=mamba_ngroups, + gdp_householder=gdp_householder, + hyena_expansion=hyena_expansion, + ) + .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) ────────────── + # ── Priming forward, then a sizing forward ──────────────────────────── + # The iteration count below is derived from a measured forward, so that + # measurement must reflect steady state. For a JIT-compiled operator it does + # not: fla's GDP kernels (and torch.compile'd FlexAttention) spend seconds + # compiling on the first call at each new shape. Sizing off that gives the + # floor of 3 iterations AND zero warmup — a JIT operator was being timed with + # 3 unwarmed samples while eager ones got 30, which is not a fair comparison. + # + # So: one untimed forward to absorb compilation, then time the second one for + # sizing. The priming pass is itself budget-checked, so a hopeless point still + # bails after one forward rather than paying for two. + torch.cuda.synchronize(device) + t0 = time.perf_counter() + with torch.inference_mode(), torch.autocast("cuda", dtype=dtype): + _ = module(x) + torch.cuda.synchronize(device) + prime_s = time.perf_counter() - t0 + + if prime_s > max_seconds: + return { + "status": "timeout", + "ms": prime_s * 1000.0, + "mem_gb": torch.cuda.max_memory_allocated(device) / (1024**3), + } + + 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.startswith("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}.") + + +# The fused 2D kernel's largest FFT tile is 128 and it requires +# max(X, Y) <= fft_size // 2, so it tops out at a 64-per-axis grid. Kept as a +# literal (not an import of fused_fftconv2d_max_spatial) so the sweep can be +# planned on a login node without subquadratic_ops_torch installed. +_FUSED_MAX_SPATIAL = 64 + + +def resolve_hyena_backend(requested: str, data_dim: int, resolution: int) -> str: + """Effective Hyena FFT backend for one point of the sweep. + + The subq_ops backends do not cover the whole (data_dim, resolution) grid, so a + single ``--fft-backend`` choice has to degrade per point rather than fail the + run. Each fallback goes to the fastest backend that *does* cover the point: + + * 3D — no CUDA FFT-conv kernel at all, so ``torch_fft``. + * ``subq_ops_fused`` in 1D/3D — no fused kernel outside 2D, so ``subq_ops`` + (which in turn becomes ``torch_fft`` in 3D). + * ``subq_ops_fused`` in 2D above ``_FUSED_MAX_SPATIAL`` — outside the fused + kernel's tile range, so ``subq_ops``. + + Returns: + The backend to construct this point's Hyena layer with. + """ + if data_dim == 3 and requested.startswith("subq_ops"): + return "torch_fft" + if requested == "subq_ops_fused": + if data_dim != 2: + return "subq_ops" + if resolution > _FUSED_MAX_SPATIAL: + return "subq_ops" + return requested + + +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.") + # Shared by mamba and gdp: upstream GDP reads the same mamba_* TransformerConfig + # fields, so a drop-in --spec swap inherits whatever the Mamba-2 layer was set to. + parser.add_argument( + "--mamba-state-dim", + type=int, + default=None, + help=( + "SSM state dim for mamba and gdp (Megatron 'mamba_state_dim'). Unset uses " + "mamba-ssm's own default for mamba, and 128 for gdp. Nemotron uses 128." + ), + ) + parser.add_argument( + "--mamba-ngroups", + type=int, + default=None, + help=( + "Key/value groups for mamba and gdp (Megatron 'mamba_num_groups'). Unset uses " + "mamba-ssm's default of 1 for mamba, and 8 for gdp. Nemotron uses 8 — leaving " + "this unset benchmarks a narrower Mamba-2 than Nemotron actually runs." + ), + ) + parser.add_argument( + "--hyena-expansion", + type=float, + default=1.0, + help=( + "Hyena inner-width multiplier 'e' (qkv h->3*e*h, mixer on e*h, out e*h->h). " + "1.0 is evo2's default, which FLOPS_MATCHING.md measures at 0.423x a Mamba-2 " + "layer -- so e=1 comparisons against mamba/gdp are NOT FLOP-matched. Use 2.36 " + "for parity with a Mamba-2 layer." + ), + ) + parser.add_argument( + "--gdp-householder", + type=int, + default=3, + help=( + "GDP householder products per token (M). 3 matches " + "gated_delta_product_original_v4.py, 2 the _nh2 variant. Hardcoded upstream " + "rather than a config field; it drives in_proj width and the kernel's L*M rows." + ), + ) + 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", "subq_ops_fused", "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( + "--short-conv", + choices=["subq_ops", "torch"], + default="subq_ops", + help=( + "Backend for Hyena's SHORT conv (distinct from --fft-backend, which is the long " + "FFT conv). 'subq_ops' uses the fused subquadratic_ops_torch.causal_conv1d kernel " + "on the causal 1D sweeps; 2D/3D are non-causal, where a causal 1D kernel does not " + "apply, so they use torch.nn.ConvNd either way." + ), + ) + 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 + # The requested backend does not cover every point; resolve it per resolution + # and report the plan up front so the coverage is visible before the sweep runs. + backend_at = {R: resolve_hyena_backend(args.fft_backend, data_dim, R) for R in resolutions} + if "hyena" in args.mixers: + for eff in dict.fromkeys(backend_at.values()): + if eff == args.fft_backend: + continue + at = [R for R in resolutions if backend_at[R] == eff] + print(f"[note] Hyena fft_backend '{args.fft_backend}' -> '{eff}' at R={at} (unsupported there).") + + # Hyena's short conv follows the operator's causality: causal (1D) gets a + # left-only-padded conv, non-causal (2D/3D) a symmetric one. The fused + # subq_ops kernel is causal and 1D, so it only applies to the former. + if data_dim == 1: + short_conv_eff = "subq_ops_causal_conv1d" if args.short_conv == "subq_ops" else "causal_conv1d_torch" + else: + short_conv_eff = f"torch_conv{data_dim}d_symmetric" + if args.short_conv == "subq_ops": + print( + f"[note] Hyena short conv: --short-conv=subq_ops needs a causal 1D operator; " + f"data_dim={data_dim} is non-causal, so using {short_conv_eff}." + ) + if "hyena" in args.mixers: + print(f"Hyena short conv: {short_conv_eff}") + + print(f"Device: {device_name}") + backend_desc = "/".join(dict.fromkeys(backend_at.values())) + 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={backend_desc} 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": backend_at[R] if mixer == "hyena" else None, + "short_conv": short_conv_eff 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=backend_at[R], + short_conv_backend=args.short_conv, + mamba_d_state=args.mamba_state_dim, + mamba_ngroups=args.mamba_ngroups, + gdp_householder=args.gdp_householder, + hyena_expansion=args.hyena_expansion, + 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": backend_at[R] if mixer == "hyena" else None, + "short_conv": short_conv_eff 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..cbf65097 100644 --- a/benchmarks/benchmark_patch_size_2d.py +++ b/benchmarks/benchmark_patch_size_2d.py @@ -106,81 +106,238 @@ 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 _short_conv_cfg( + hidden_dim: int, + *, + data_dim: int, + is_causal: bool, + backend: str = "subq_ops", + kernel_size: int = 3, + inner_dim: int | None = None, +) -> LazyConfig: + """Build HyenaND's short-conv config, matching the operator's causality. + + The short conv follows ``is_causal`` rather than always using symmetric + padding. Previously the causal 1D config paired a causal long conv with a + ``padding=1`` symmetric short conv, which let the operator see one token of + future context — causal in the long conv only. + + * ``is_causal=True`` (1D): left-only padding. With ``backend="subq_ops"`` + this is :class:`SubqOpsCausalConv1d`, the fused + ``subquadratic_ops_torch.causal_conv1d`` CUDA kernel — depthwise-only + (satisfied here: groups == in_channels == out_channels == 3*hidden_dim), + stride/dilation 1, 1D only. Otherwise :class:`CausalConv1D`, the portable + ``F.conv1d`` equivalent. + * ``is_causal=False`` (2D/3D): symmetric-padded ``torch.nn.ConvNd``, where + the full context is meant to be visible. The fused kernel is causal and 1D, + so it does not apply; ``backend="subq_ops"`` falls back here. + + Returns: + The short-conv ``LazyConfig``. + + Raises: + ValueError: If ``backend`` is not ``"torch"`` or ``"subq_ops"``. + """ + # The short conv runs on the mixer's INNER width (3 * e * hidden_dim when + # expanded), not on hidden_dim. + hidden_dim = hidden_dim if inner_dim is None else inner_dim + if backend not in ("torch", "subq_ops"): + raise ValueError(f"short_conv backend must be 'torch' or 'subq_ops'. Got {backend!r}.") + + if not is_causal: + return LazyConfig(_CONV_ND[data_dim])( + in_channels=3 * hidden_dim, + out_channels=3 * hidden_dim, + kernel_size=kernel_size, + groups=3 * hidden_dim, + padding=kernel_size // 2, + bias=False, + ) + + if backend == "subq_ops" and data_dim == 1: + from nvsubquadratic.modules.subq_ops_causal_conv1d import SubqOpsCausalConv1d + + return LazyConfig(SubqOpsCausalConv1d)( + in_channels=3 * hidden_dim, + out_channels=3 * hidden_dim, + kernel_size=kernel_size, + groups=3 * hidden_dim, + bias=False, + ) + + from nvsubquadratic.modules.causal_conv1d import CausalConv1D + + return LazyConfig(CausalConv1D)( + in_channels=3 * hidden_dim, + out_channels=3 * hidden_dim, + kernel_size=kernel_size, + groups=3 * hidden_dim, + bias=False, + ) + + +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, + short_conv_backend: str = "subq_ops", + expansion: float = 1.0, +) -> LazyConfig: + """Build the HyenaND mixer. + + ``expansion`` is the ``e`` of the FLOP-matching analysis: the mixer runs at an + inner width ``e * hidden_dim`` (``qkv_proj: h -> 3*e*h``, long conv and short conv + on ``e*h`` channels, ``out_proj: e*h -> h``). ``e = 1`` is evo2's default, which + FLOPS_MATCHING.md measures at only **0.423x a Mamba-2 layer** — so comparisons + against Mamba-2 or GDP at ``e = 1`` are NOT FLOP-matched and understate Hyena's + cost. That document derives ``e = 2.36`` for parity. + """ + inner_dim = round(expansion * hidden_dim) return LazyConfig(QKVSequenceMixer)( hidden_dim=hidden_dim, + inner_dim=inner_dim, mixer_cfg=LazyConfig(Hyena)( global_conv_cfg=LazyConfig(CKConvND)( - data_dim=DATA_DIM, - hidden_dim=hidden_dim, + data_dim=data_dim, + hidden_dim=inner_dim, kernel_cfg=LazyConfig(SIRENKernelND)( - data_dim=DATA_DIM, - out_dim=hidden_dim, + data_dim=data_dim, + out_dim=inner_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)( - in_channels=3 * hidden_dim, - out_channels=3 * hidden_dim, - kernel_size=3, - groups=3 * hidden_dim, - padding=1, - bias=False, + short_conv_cfg=_short_conv_cfg( + hidden_dim, + data_dim=data_dim, + is_causal=is_causal, + backend=short_conv_backend, + inner_dim=inner_dim, ), gate_nonlinear_cfg=LazyConfig(torch.nn.Identity)(), - pixelhyena_norm_cfg=LazyConfig(torch.nn.LayerNorm)(normalized_shape=hidden_dim), + pixelhyena_norm_cfg=LazyConfig(torch.nn.LayerNorm)(normalized_shape=inner_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, + d_state: int | None = None, + ngroups: int | None = None, +) -> LazyConfig: + """Build the Mamba2 baseline mixer. + + ``d_state``/``ngroups`` default to mamba-ssm's own defaults when ``None``. Pass them + to match a specific deployment: Nemotron runs ``state_dim 128, num_groups 8``, and + mamba-ssm defaults ``ngroups=1``, so leaving it unset silently benchmarks a narrower + layer than Nemotron's. + """ + # 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 + mamba_kwargs = {"d_model": hidden_dim, "headdim": headdim, "expand": expand} + if d_state is not None: + mamba_kwargs["d_state"] = d_state + if ngroups is not None: + mamba_kwargs["ngroups"] = ngroups + return LazyConfig(MambaNDMixer)( - mamba_layer_cfg=LazyConfig(Mamba2)( + mamba_layer_cfg=LazyConfig(Mamba2)(**mamba_kwargs), + bidirectional=bidirectional, + ) + + +def _gdp_mixer_cfg( + hidden_dim: int, + *, + num_householder: int = 3, + d_state: int = 128, + headdim: int = 64, + ngroups: int = 8, + nheads: int | None = None, + expand: int = 2, +) -> LazyConfig: + """Build the Gated Delta Product baseline — the mixer Nemotron may adopt over Mamba-2. + + Defaults mirror ``TransformerConfig``'s ``mamba_*`` fields, which is what GDP reads: + ``mamba_state_dim=128``, ``mamba_head_dim=64``, ``mamba_num_groups=8``, and + ``mamba_num_heads`` falling back to ``hidden*expand // head_dim``. ``num_householder`` + is NOT a config field upstream — it is hardcoded per source file (3 in + ``gated_delta_product_original_v4.py``, 2 in ``_nh2``) — so a drop-in ``--spec`` swap + inherits the Mamba-2 config and widens ``in_proj`` by the householder count. + """ + from gated_delta_product_ref import GatedDeltaProduct, GatedDeltaProductNDMixer + + return LazyConfig(GatedDeltaProductNDMixer)( + mixer_cfg=LazyConfig(GatedDeltaProduct)( d_model=hidden_dim, - headdim=MAMBA_HEADDIM, - expand=MAMBA_EXPAND, - ), - bidirectional=MAMBA_BIDIRECTIONAL, + num_householder=num_householder, + d_state=d_state, + headdim=headdim, + ngroups=ngroups, + nheads=nheads, + expand=expand, + ) ) diff --git a/benchmarks/gated_delta_product_ref.py b/benchmarks/gated_delta_product_ref.py new file mode 100644 index 00000000..21567189 --- /dev/null +++ b/benchmarks/gated_delta_product_ref.py @@ -0,0 +1,300 @@ +# 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. + +"""Gated Delta Product (GDP) — a standalone port of Megatron's mixer, for benchmarking. + +GDP is the mixer Nemotron is evaluating as a replacement for its Mamba-2 layers, so it +is the operator HyenaND actually has to beat. This is a **benchmark baseline**, not an +nvSubquadratic contribution: it exists so ``gdp`` can appear alongside ``mamba`` in the +forward-time sweeps. + +Ported from, and kept faithful to:: + + megatron-lm @ dev-arch-mar2026-v2 (daf96b8e) + megatron/core/ssm/gated_delta_product_original_v4.py (num_householder=3) + megatron/core/ssm/gated_delta_product_original_v4_nh2.py (num_householder=2) + +The two upstream files differ only in ``num_householder``; here that is a constructor +argument, defaulting to **3** to match ``gated_delta_product_original_v4.py`` — the +actively maintained file, and the one nemotron_workspace's HANDOFF cites as the +reference. Pass ``num_householder=2`` for the ``_nh2`` variant. + +**Why a port rather than importing Megatron.** The operator being measured is +``fla.ops.gated_delta_product.chunk_gated_delta_product``, an external library kernel +that is identical either way. Importing the Megatron module instead would require +TransformerEngine (its spec builds ``in_proj``/``out_proj`` from +``TELayerNormColumnParallelLinear``/``TERowParallelLinear``), a megatron-core built from +the ADLR branch (the released 0.18.2 has no GDP at all), and an initialised +``ProcessGroupCollection`` — all of which are inert at TP=1/CP=1 on one GPU. + +**What is deliberately dropped**, none of which does arithmetic at TP=1/CP=1 on a single +device: tensor/context parallelism (``GDPContextParallel``), inference contexts (static +and dynamic batching, conv/ssm state caches), sequence packing (``cu_seqlens``, +``seq_idx``), and sharded-checkpoint plumbing. + +**What differs numerically from upstream**: ``in_proj``/``out_proj`` are +``torch.nn.Linear`` rather than TE fused linears. At TP=1 both are plain GEMMs, but TE's +fused epilogues can be faster, so absolute timings here are a slight *under*-estimate of +GDP's projection speed and a correspondingly conservative comparison for HyenaND. The +FLA kernel — which dominates at long sequence — is the same code upstream calls. + +Requires ``flash-linear-attention`` (``fla``), ``mamba-ssm`` (for ``RMSNormGated`` and +``causal_conv1d``) and ``einops``. +""" + +from __future__ import annotations + +import math +import sys + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def _shim_tilelang() -> None: + """Make ``import tilelang`` fail cleanly, before anything imports ``mamba_ssm``. + + ``mamba_ssm``'s package ``__init__`` eagerly imports ``modules.mamba3`` -> + ``tilelang`` -> a bundled TVM that collides with the ``apache-tvm-ffi`` that + subquadratic-ops-torch >= 0.2.2 installs, raising + ``AttributeError: attribute '__dict__' of 'type' objects is not writable``. + Setting the module to ``None`` turns that into an ImportError, which mamba_ssm's + optional import skips — the same shim ``benchmark_patch_size_2d._mamba_mixer_cfg`` + uses for Mamba2. + + This must run before the FIRST ``mamba_ssm`` or ``fla`` import in the process, not + merely before the ``fla`` one: importing *any* ``mamba_ssm`` submodule (e.g. + ``ops.triton.layernorm_gated`` for ``RMSNormGated``) executes that package + ``__init__`` and trips the same chain. + """ + sys.modules.setdefault("tilelang", None) + + +def _import_fla(): + """Return the FLA chunked GDP kernel.""" + _shim_tilelang() + from fla.ops.gated_delta_product import chunk_gated_delta_product + + return chunk_gated_delta_product + + +class GatedDeltaProduct(nn.Module): + """Single GDP mixer layer, ``[B, L, D] -> [B, L, D]``. + + Upstream is sequence-first (``[L, B, D]``); this takes batch-first so it drops into + the benchmark harness beside the other mixers, and transposes internally. + + Args: + d_model: Model hidden size. + num_householder: Householder products per token (``M``). 3 matches + ``gated_delta_product_original_v4.py``, 2 the ``_nh2`` variant. Drives both + the ``in_proj`` width and the kernel's effective sequence length (``L*M``), + so it is the dominant cost knob. + d_state: SSM state dim (Megatron ``mamba_state_dim``). + headdim: Per-head width (Megatron ``mamba_head_dim``). + ngroups: Key/query groups (Megatron ``mamba_num_groups``). When + ``nheads > ngroups`` the kernel's q/k are repeat-interleaved, GQA-style. + nheads: Value heads. Defaults to ``d_model * expand // headdim``, matching + ``TransformerConfig.mamba_num_heads``'s documented fallback. + expand: Inner expansion factor, used only for the ``nheads`` default. + d_conv: Causal short-conv width. + conv_bias / bias: Bias on the short conv / the projections. + rmsnorm: Apply the gated RMSNorm before ``out_proj``. + dt_min, dt_max, dt_init_floor: ``dt_bias`` init range. + A_init_range: ``A`` init range; ``A_log = log(A)`` is kept in fp32. + device, dtype: Passed to the submodules. + """ + + def __init__( + self, + d_model: int, + *, + num_householder: int = 3, + d_state: int = 128, + headdim: int = 64, + ngroups: int = 8, + nheads: int | None = None, + expand: int = 2, + d_conv: int = 4, + conv_bias: bool = False, + bias: bool = False, + rmsnorm: bool = True, + norm_before_gate: bool = False, + dt_min: float = 0.001, + dt_max: float = 0.1, + dt_init_floor: float = 1e-4, + A_init_range: tuple[float, float] = (0.0, 16.0), + conv_init: float | None = None, + device=None, + dtype=None, + ) -> None: + super().__init__() + if num_householder < 1: + raise ValueError(f"num_householder must be >= 1, got {num_householder}.") + + self.d_model = d_model + self.num_householder = num_householder + self.d_state = d_state + self.headdim = headdim + self.ngroups = ngroups + self.nheads = nheads if nheads is not None else (d_model * expand) // headdim + self.d_inner = self.nheads * self.headdim + self.d_conv = d_conv + self.rmsnorm = rmsnorm + self.activation = "silu" + + if self.nheads % self.ngroups != 0: + raise ValueError( + f"nheads ({self.nheads}) must be divisible by ngroups ({self.ngroups}) for the q/k repeat_interleave." + ) + + M = num_householder + factory = {"device": device, "dtype": dtype} + + # z | V (M copies) | K (M copies) | Q | b (M per head) | a (1 per head). + # Width formula copied verbatim from the upstream in_proj construction. + in_width = self.d_inner * (1 + M) + self.ngroups * self.d_state * (M + 1) + self.nheads * (M + 1) + self.in_proj = nn.Linear(d_model, in_width, bias=bias, **factory) + + # Short conv spans V and K (M copies each) plus one Q — not z, and not b/a. + conv_dim = self.d_inner * M + (M + 1) * self.ngroups * self.d_state + self.conv_dim = conv_dim + self.conv1d = nn.Conv1d( + in_channels=conv_dim, + out_channels=conv_dim, + bias=conv_bias, + kernel_size=d_conv, + groups=conv_dim, + padding=d_conv - 1, + **factory, + ) + if conv_init is not None: + nn.init.uniform_(self.conv1d.weight, -conv_init, conv_init) + + # dt_bias init so softplus(dt_bias) lands in [dt_min, dt_max]; A_log stays fp32 + # or A can underflow to -inf in fp16 (upstream's comment). + dt = torch.exp( + torch.rand(self.nheads, device=device) * (math.log(dt_max) - math.log(dt_min)) + math.log(dt_min) + ).clamp(min=dt_init_floor) + self.dt_bias = nn.Parameter(dt + torch.log(-torch.expm1(-dt))) + A = torch.empty(self.nheads, dtype=torch.float32, device=device).uniform_(*A_init_range) + self.A_log = nn.Parameter(torch.log(A)) + + if rmsnorm: + _shim_tilelang() # must precede the first mamba_ssm import in the process + from mamba_ssm.ops.triton.layernorm_gated import RMSNorm as RMSNormGated + + self.norm = RMSNormGated( + self.d_inner, + eps=1e-5, + group_size=self.d_inner // self.ngroups, + norm_before_gate=norm_before_gate, + **factory, + ) + + self.out_proj = nn.Linear(self.d_inner, d_model, bias=bias, **factory) + + self._chunk_gated_delta_product = None + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Run the mixer. ``hidden_states`` is ``[B, L, D]``; returns the same shape.""" + from causal_conv1d import causal_conv1d_fn + from einops import rearrange + + if self._chunk_gated_delta_product is None: + self._chunk_gated_delta_product = _import_fla() + chunk_gated_delta_product = self._chunk_gated_delta_product + + M, H, P, G, S = self.num_householder, self.nheads, self.headdim, self.ngroups, self.d_state + + zVKQba = self.in_proj(hidden_states) + z, VKQ, ba = torch.split( + zVKQba, + [self.d_inner, self.conv_dim, H * (M + 1)], + dim=-1, + ) + + # causal_conv1d_fn wants [B, D, L], channels-first contiguous. + VKQ = rearrange(VKQ, "b l d -> b d l").contiguous() + VKQ = causal_conv1d_fn( + x=VKQ, + weight=rearrange(self.conv1d.weight, "d 1 w -> d w"), + bias=self.conv1d.bias, + activation=self.activation, + ) + VKQ = rearrange(VKQ, "b d l -> b l d").contiguous() + + value, key, query = torch.split(VKQ, [self.d_inner * M, G * S * M, G * S], dim=-1) + b, a = torch.split(ba, [H * M, H], dim=-1) + + # The M householder copies are folded into the sequence axis: the kernel sees + # L*M rows for v/k/beta while q and g stay at L, one query per real token. + z = rearrange(z, "b l (h p) -> b l h p", p=P).contiguous() + value = rearrange(value, "b l (m h p) -> b (l m) h p", m=M, p=P).contiguous() + key = rearrange(key, "b l (m g n) -> b (l m) g n", m=M, n=S).contiguous() + query = rearrange(query, "b l (g n) -> b l g n", n=S).contiguous() + beta = rearrange(b.contiguous().sigmoid(), "b l (m h) -> b (l m) h", m=M).contiguous() + + # fp32 for A: in fp16 exp(A_log) can underflow and give -inf decay. + g = -self.A_log.float().exp() * F.softplus(a.float().contiguous() + self.dt_bias) + + if H // G > 1: + query = query.repeat_interleave(H // G, dim=2) + key = key.repeat_interleave(H // G, dim=2) + + core_attn_out, _ = chunk_gated_delta_product( + query, + key, + value, + g=g, + beta=beta, + initial_state=None, + output_final_state=False, + num_householder=M, + use_qk_l2norm_in_kernel=True, + ) + + y = rearrange(core_attn_out, "b l h p -> b l (h p)").contiguous() + if self.rmsnorm: + y = self.norm(y, rearrange(z, "b l h p -> b l (h p)").contiguous()) + return self.out_proj(y) + + +class GatedDeltaProductNDMixer(nn.Module): + """Adapter giving :class:`GatedDeltaProduct` the benchmark's ND tensor convention. + + The harness feeds ``[B, R, ...(N axes)..., C]``; GDP is a causal 1D sequence mixer, + so the spatial axes are flattened to one scan and restored afterwards — the same + rasterisation the Mamba2 baseline uses, which is what makes the two comparable. + + Args: + mixer_cfg: A ``LazyConfig`` for the inner :class:`GatedDeltaProduct`, following + the repo's nested-config convention (as ``MambaNDMixer`` takes + ``mamba_layer_cfg``): the config is instantiated here rather than by the + caller, so ``instantiate`` does not have to recurse into the kwarg. + """ + + def __init__(self, mixer_cfg) -> None: + super().__init__() + from nvsubquadratic.lazy_config import instantiate + + self.mixer = instantiate(mixer_cfg) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Flatten the spatial axes, mix, restore. ``x`` is ``[B, ...spatial, C]``.""" + shape = x.shape + y = self.mixer(x.reshape(shape[0], -1, shape[-1])) + return y.reshape(shape) diff --git a/benchmarks/results/.gitignore b/benchmarks/results/.gitignore new file mode 100644 index 00000000..056d5640 --- /dev/null +++ b/benchmarks/results/.gitignore @@ -0,0 +1,14 @@ +# Slurm job logs — one pair per run, large and specific to a job id. +*.out + +# Plots — regenerable from the JSONL beside them: +# PYTHONPATH=. python scripts/visualization/visualize_forward_time_nd.py \ +# --input .jsonl --out .png +*.png +*.pdf + +# The JSONL rows and BENCHMARK_RESULTS.md are the primary record and ARE tracked. + +# Smoke / bring-up runs — superseded by the production sweeps, and some contain +# rows from since-fixed harness bugs. Kept on disk, out of the record. +*smoke* diff --git a/benchmarks/results/BENCHMARK_RESULTS.md b/benchmarks/results/BENCHMARK_RESULTS.md new file mode 100644 index 00000000..6ccde980 --- /dev/null +++ b/benchmarks/results/BENCHMARK_RESULTS.md @@ -0,0 +1,383 @@ +# nvSubquadratic forward-time benchmarks — results digest + +Single-layer forward-pass timings for HyenaND vs attention kernels vs Mamba2, +swept from a 16-wide grid up to ~16.7M tokens in 1D, 2D and 3D. + +This file is generated from the JSONL sweep outputs in this directory; every number +below is read from that data rather than transcribed. Raw per-point records live in +`forward_time_{1,2,3}d.jsonl` and `forward_time_flash_{1,2,3}d.jsonl`, with matching +`.png`/`.pdf` plots and `*-.out` run logs. + +______________________________________________________________________ + +## What was run + +Two sweep configurations, each across 1D/2D/3D — six jobs total. They answer +different questions and are **not** comparable to each other: + +| Config | hidden_dim | head_dim | Mixers | Reaches | Purpose | +| ---------------- | ---------- | -------- | ---------------------------------- | ------------ | ------------------------------------------------------------------------------------------------------------------------------------------ | +| **Reach** | 8 | 4 | hyena, attention, mamba | 16.7M tokens | How far each operator scales. Small width keeps the qkv tensor under torch's 2^31 index limit at 16.7M. | +| **Flash-kernel** | 512 | 128 | hyena, attention, flex, fa4, mamba | ~1M tokens | SDPA vs FlexAttention vs FlashAttention-4 vs HyenaND at the width flash kernels are optimised for. Walls at ~1.4M on the same index limit. | + +`attention`, `flex` and `fa4` are three interchangeable kernels on one shared q/k/v + +RoPE path. `flex` and `fa4` require head_dim >= 16, so they cannot run in the Reach +config and appear only in the Flash-kernel sweeps. + +______________________________________________________________________ + +## Environment + +| Component | Version | +| --------------------------- | ---------------------------------------------------------------------------- | +| GPU | NVIDIA GB200 (aarch64), driver 580.173.02 | +| torch | 2.12.1+cu130 (CUDA 13.0) | +| subquadratic-ops-torch-cu13 | 0.2.2 | +| mamba-ssm / causal-conv1d | 2.3.2.post1 / 1.6.2.post1 | +| flash-attn-4 | 4.0.0b23 (cutlass-dsl 4.6.0.dev0) | +| apex | 0.1 | +| Image | `/lustre/fsw/healthcareeng_bionemo/farhadr/enroot/nvsubquadratic-arm64.sqsh` | +| Branch | `farhadr/2d_bench` (merge of PR #137 + PR #138) | + +______________________________________________________________________ + +## Results + +### Reach 1D (`forward_time_1d`) + +- SLURM job: `2559702` | device: `NVIDIA GB200` | dtype: `bf16` +- hidden_dim=8, num_heads=2, batch_size=1, data_dim=1 (L = R^1) +- Hyena short conv: subq_ops_causal_conv1d +- Hyena FFT backend: `subq_ops` at R=16..16777216 + +| R | L | hyena | attention | mamba | +| ---------- | ---------- | ------ | ---------- | ------ | +| 16 | 16 | 1.667 | 0.652 | 2.237 | +| 32 | 32 | 1.508 | 1.144 | 2.410 | +| 64 | 64 | 1.515 | 1.143 | 2.408 | +| 128 | 128 | 1.470 | 1.108 | 2.411 | +| 256 | 256 | 1.431 | 1.149 | 2.412 | +| 512 | 512 | 1.412 | 1.134 | 2.272 | +| 1,024 | 1,024 | 1.453 | 1.178 | 2.391 | +| 2,048 | 2,048 | 1.450 | 0.643 | 2.415 | +| 4,096 | 4,096 | 1.452 | 1.188 | 2.279 | +| 8,192 | 8,192 | 1.453 | 1.139 | 2.425 | +| 16,384 | 16,384 | 1.405 | 1.092 | 2.445 | +| 32,768 | 32,768 | 1.420 | 1.395 | 2.439 | +| 65,536 | 65,536 | 1.497 | 5.024 | 2.437 | +| 131,072 | 131,072 | 1.501 | 17.164 | 2.446 | +| 262,144 | 262,144 | 1.480 | 67.812 | 3.067 | +| 524,288 | 524,288 | 1.921 | 261.968 | 5.956 | +| 1,048,576 | 1,048,576 | 4.348 | 1046.639 | 11.772 | +| 2,097,152 | 2,097,152 | 8.865 | 4185.474 | 23.352 | +| 4,194,304 | 4,194,304 | 18.639 | 16788.312 | error | +| 8,388,608 | 8,388,608 | 37.347 | 67089.927 | error | +| 16,777,216 | 16,777,216 | 74.507 | 268037.938 | error | + +All values are ms per forward pass, lower is better. + +- max attention/hyena: **3597x** at L=16,777,216 +- max mamba/hyena: **3.1x** at L=524,288 + +Points without a timing: + +- `mamba` error at R=\[4194304, 8388608, 16777216\] + +______________________________________________________________________ + +### Reach 2D (`forward_time_2d`) + +- SLURM job: `2559620` | device: `NVIDIA GB200` | dtype: `bf16` +- hidden_dim=8, num_heads=2, batch_size=1, data_dim=2 (L = R^2) +- Hyena short conv: torch_conv2d_symmetric +- Hyena FFT backend: `subq_ops` at R=128..4096; `subq_ops_fused` at R=16..64 + +| R | L | hyena | attention | mamba | +| ----- | ---------- | ------ | ---------- | ------ | +| 16 | 256 | 1.299 | 0.803 | 2.289 | +| 32 | 1,024 | 1.113 | 1.407 | 2.154 | +| 64 | 4,096 | 1.093 | 1.483 | 2.177 | +| 128 | 16,384 | 1.143 | 1.443 | 2.332 | +| 256 | 65,536 | 1.046 | 5.096 | 2.338 | +| 512 | 262,144 | 1.055 | 67.892 | 3.198 | +| 1,024 | 1,048,576 | 2.361 | 1046.963 | 12.224 | +| 2,048 | 4,194,304 | 10.156 | 16787.928 | error | +| 4,096 | 16,777,216 | 41.292 | 268220.812 | error | + +All values are ms per forward pass, lower is better. + +- max attention/hyena: **6496x** at L=16,777,216 +- max mamba/hyena: **5.2x** at L=1,048,576 + +Points without a timing: + +- `mamba` error at R=\[2048, 4096\] + +______________________________________________________________________ + +### Reach 3D (`forward_time_3d`) + +- SLURM job: `2559621` | device: `NVIDIA GB200` | dtype: `bf16` +- hidden_dim=8, num_heads=2, batch_size=1, data_dim=3 (L = R^3) +- Hyena short conv: torch_conv3d_symmetric +- Hyena FFT backend: `torch_fft` at R=16..256 + +| R | L | hyena | attention | mamba | +| --- | ---------- | ------ | ---------- | ------ | +| 16 | 4,096 | 1.118 | 0.579 | 2.401 | +| 32 | 32,768 | 1.115 | 1.334 | 2.414 | +| 64 | 262,144 | 1.122 | 67.674 | 3.142 | +| 128 | 2,097,152 | 7.579 | 4178.602 | 23.841 | +| 256 | 16,777,216 | 62.031 | 267676.531 | error | + +All values are ms per forward pass, lower is better. + +- max attention/hyena: **4315x** at L=16,777,216 +- max mamba/hyena: **3.1x** at L=2,097,152 + +Points without a timing: + +- `mamba` error at R=\[256\] + +______________________________________________________________________ + +### Flash-kernel 1D (`forward_time_flash_1d`) + +- SLURM job: `2559703` | device: `NVIDIA GB200` | dtype: `bf16` +- hidden_dim=512, num_heads=4, batch_size=1, data_dim=1 (L = R^1) +- Hyena short conv: subq_ops_causal_conv1d +- Hyena FFT backend: `subq_ops` at R=16..1048576 + +| R | L | hyena | attention | flex | fa4 | mamba | +| --------- | --------- | ------ | --------- | -------- | -------- | ------ | +| 16 | 16 | 1.956 | 0.551 | 0.997 | 0.825 | 2.286 | +| 32 | 32 | 1.428 | 0.558 | 0.893 | 0.950 | 2.291 | +| 64 | 64 | 1.366 | 0.562 | 1.069 | 0.949 | 2.299 | +| 128 | 128 | 1.332 | 0.548 | 0.981 | 0.955 | 2.321 | +| 256 | 256 | 1.328 | 0.617 | 1.024 | 0.759 | 2.297 | +| 512 | 512 | 1.322 | 0.617 | 1.049 | 0.782 | 2.187 | +| 1,024 | 1,024 | 1.326 | 0.614 | 1.043 | 0.938 | 2.331 | +| 2,048 | 2,048 | 1.289 | 0.611 | 1.036 | 0.961 | 2.358 | +| 4,096 | 4,096 | 1.286 | 0.632 | 1.046 | 0.960 | 2.190 | +| 8,192 | 8,192 | 1.292 | 0.631 | 1.056 | 0.959 | 2.297 | +| 16,384 | 16,384 | 1.290 | 0.822 | 2.425 | 0.996 | 2.322 | +| 32,768 | 32,768 | 1.286 | 2.228 | 7.597 | 2.658 | 2.310 | +| 65,536 | 65,536 | 2.395 | 6.965 | 28.295 | 8.120 | 3.007 | +| 131,072 | 131,072 | 4.761 | 24.981 | 114.488 | 30.295 | 5.819 | +| 262,144 | 262,144 | 9.971 | 95.353 | 460.909 | 119.907 | 11.451 | +| 524,288 | 524,288 | 24.203 | 379.858 | 1831.237 | 505.537 | 22.792 | +| 1,048,576 | 1,048,576 | 76.215 | 1587.445 | 7299.468 | 2085.962 | - | + +All values are ms per forward pass, lower is better. + +- max attention/hyena: **21x** at L=1,048,576 +- max flex/hyena: **96x** at L=1,048,576 +- max fa4/hyena: **27x** at L=1,048,576 +- max mamba/hyena: **1.8x** at L=2,048 + +______________________________________________________________________ + +### Flash-kernel 2D (`forward_time_flash_2d`) + +- SLURM job: `2559623` | device: `NVIDIA GB200` | dtype: `bf16` +- hidden_dim=512, num_heads=4, batch_size=1, data_dim=2 (L = R^2) +- Hyena short conv: torch_conv2d_symmetric +- Hyena FFT backend: `subq_ops` at R=128..1024; `subq_ops_fused` at R=16..64 + +| R | L | hyena | attention | flex | fa4 | mamba | +| ----- | --------- | ------ | --------- | -------- | -------- | ----- | +| 16 | 256 | 1.625 | 0.986 | 1.242 | 1.065 | 2.562 | +| 32 | 1,024 | 1.226 | 0.949 | 1.391 | 1.037 | 2.511 | +| 64 | 4,096 | 1.211 | 0.966 | 1.544 | 1.448 | 2.551 | +| 128 | 16,384 | 1.359 | 0.999 | 2.609 | 1.458 | 2.736 | +| 256 | 65,536 | 3.555 | 7.263 | 28.950 | 8.412 | 2.725 | +| 512 | 262,144 | 14.075 | 96.786 | 467.332 | 122.081 | 8.659 | +| 1,024 | 1,048,576 | 89.187 | 1648.564 | 7381.654 | 2080.850 | - | + +All values are ms per forward pass, lower is better. + +- max attention/hyena: **18x** at L=1,048,576 +- max flex/hyena: **83x** at L=1,048,576 +- max fa4/hyena: **23x** at L=1,048,576 +- max mamba/hyena: **2.1x** at L=4,096 + +______________________________________________________________________ + +### Flash-kernel 3D (`forward_time_flash_3d`) + +- SLURM job: `2559624` | device: `NVIDIA GB200` | dtype: `bf16` +- hidden_dim=512, num_heads=4, batch_size=1, data_dim=3 (L = R^3) +- Hyena short conv: torch_conv3d_symmetric +- Hyena FFT backend: `torch_fft` at R=16..64 + +| R | L | hyena | attention | flex | fa4 | mamba | +| --- | ------- | ------ | --------- | ------- | ------- | ----- | +| 16 | 4,096 | 1.136 | 0.496 | 0.730 | 0.708 | 2.707 | +| 32 | 32,768 | 7.897 | 1.764 | 7.329 | 2.150 | 2.929 | +| 64 | 262,144 | 70.852 | 96.806 | 459.914 | 118.767 | 8.324 | + +All values are ms per forward pass, lower is better. + +- max attention/hyena: **1.4x** at L=262,144 +- max flex/hyena: **6.5x** at L=262,144 +- max fa4/hyena: **1.7x** at L=262,144 +- max mamba/hyena: **2.4x** at L=4,096 + +______________________________________________________________________ + +## How to read these numbers + +### 1. Below ~65K tokens the sweep does not measure attention + +In the Reach config (head_dim 4) attention sits flat at ~1.1 ms from R=32 to R=32768. +That floor is fixed per-iteration dispatch overhead, not attention: the quadratic term +is real but far below the floor until L=65536, where attention first departs (5.02 ms) +and then scales cleanly (17.2 -> 67.8 -> 262 -> 1047 ms, ~4x per doubling). + +**Do not draw conclusions from points below L=65536.** For a figure, truncating the +x-axis there is the honest presentation. + +### 2. There is a reproducible dip at R=2048 in the 1D Reach plot + +Investigated directly. It is **not** hardware and **not** a one-off: + +- Clean ECC (0 uncorrected, all 4 GPUs), no throttling of any kind, 27 C. +- Reproduces on a different node (`ptyche0178`) than the production run (`ptyche0077`), + ruling out a node-specific fault. +- 6 repetitions at R=2048: min 0.585, median 0.649, max 1.189 ms. Five of six land at + ~0.65 ms, so the production value of 0.643 is the *typical* reading, not an outlier. + +Timings in this region are **bimodal**, landing near ~0.6 ms or ~1.2 ms, and which +resolutions land where changes with node and process context (in production only R=16 +and R=2048 were fast; on the rerun node R=1024 and R=2048 were fast while R=4096 was +slow). The likely mechanism is SDPA backend selection: at head_dim 4 the flash path is +ineligible, so PyTorch chooses between math and mem-efficient kernels by heuristic, +giving two stable states. Re-running relocates the dip rather than removing it. + +### 3. Mamba's `error` points are kernel limits, NOT out-of-memory + +No point in any sweep ran out of memory. The `error` entries here are +`cudaErrorInvalidValue` raised by Dao-AILab's `causal_conv1d`: its channels-last path +tiles the sequence at 64 elements per block and hits the CUDA grid-dimension cap of +65,535, giving a hard ceiling of **4,194,240 tokens**. That is a launch-configuration +limit, so it occurs on any GPU regardless of memory (peak use here was under 5 GB). + +At wider configurations a *different* and earlier limit binds first: 32-bit index +overflow in the Triton SSD scan (`ssd_chunk_state.py`), which fails at ~2^31 elements +and therefore moves inversely with model width (~600K tokens at hidden 768). These +sweeps use hidden_dim 8, which is why they reach 4M before failing. + +Both limits are implementation properties of the shipping kernels, not architectural +bounds on Mamba-2, and both are documented with reproductions in +`docs/mamba2_limits.md`. The Mamba baseline intentionally keeps its own +`causal_conv1d`; it was not switched to the subq_ops kernel. + +### 4. The fused 2D kernel covers only part of the 2D sweep + +`fft_backend=subq_ops_fused` (subq_ops >= 0.2.2) is **2D-only and capped at 64 per +axis** — its largest FFT tile is 128 and it requires max(X,Y) \<= fft_size/2. Across a +16..16M sweep it therefore covers 2D at R=16/32/64 only; R>=128 falls back to +`subq_ops`, and 3D to `torch_fft` (no 3D CUDA kernel exists). The backend actually used +is recorded per row in the JSONL as `backend`, so no point is misattributed. + +### 5. 1D HyenaND is not comparable to runs before 2026-08-10 + +HyenaND's short conv now follows the operator's causality. Previously the causal 1D +config paired a causal long conv with a symmetric-padded short conv, so the operator +could see one token of future context. 1D now uses the fused +`subquadratic_ops_torch.causal_conv1d` (left-only padding), recorded per row as +`short_conv`. This is a behaviour change, not only a speedup. 2D/3D are unaffected: +they are non-causal and still use symmetric `torch.nn.ConvNd`, because the fused kernel +is causal and 1D-only. + +### 6. The Flash-kernel 3D sweep is too short to show a crossover + +At hidden_dim 512 in 3D the 2^31 index limit bites at R=64 (262K tokens), leaving only +three points. HyenaND is still *behind* attention at the first two and only pulls ahead +at the last (70.9 vs 96.8 ms, 1.4x). Treat that sweep as a kernel-availability check +rather than evidence about scaling; the 3D scaling story is in the Reach 3D sweep, +which reaches 16.7M. + +### 7. Points in the tail are averaged over very few iterations + +Each datapoint is a single measurement: CUDA events around one back-to-back loop of N +forwards, divided by N. It is **not** an average of repeated runs. N adapts to the cost +(target ~5 s per point, capped by `--num-iters` and `--max-seconds-per-point`), so the +largest points have very few iterations — the 16.7M attention point is N=1. Those are +still reliable because a 268-second measurement swamps jitter, but they are single-shot. + +1D attention iteration counts: + +| L | iterations | +| ---------- | ---------- | +| 16 | 30 | +| 32 | 30 | +| 64 | 30 | +| 128 | 30 | +| 256 | 30 | +| 512 | 30 | +| 1,024 | 30 | +| 2,048 | 30 | +| 4,096 | 30 | +| 8,192 | 30 | +| 16,384 | 30 | +| 32,768 | 30 | +| 65,536 | 30 | +| 131,072 | 30 | +| 262,144 | 30 | +| 524,288 | 19 | +| 1,048,576 | 5 | +| 2,097,152 | 3 | +| 4,194,304 | 3 | +| 8,388,608 | 3 | +| 16,777,216 | 1 | + +______________________________________________________________________ + +## JSONL schema + +One record per (mixer, resolution): + +```json +{ + "mixer": "attention", + "resolution": 16, + "seq_len": 16, + "data_dim": 1, + "backend": null, + "short_conv": null, + "batch_size": 1, + "hidden_dim": 8, + "num_heads": 2, + "dtype": "bf16", + "device": "NVIDIA GB200", + "status": "ok", + "ms": 0.6519936243693034, + "mem_gb": 0.03125905990600586, + "iters": 30 +} +``` + +| Field | Meaning | +| ------------ | ------------------------------------------------------------------------------------------------- | +| `status` | `ok`, `oom`, `timeout`, `error`, or `unavailable` (dependency missing). Only `ok` rows have `ms`. | +| `ms` | Milliseconds per forward pass. | +| `iters` | Iterations the timed loop averaged over (adaptive; see note 6). | +| `backend` | Hyena FFT-conv backend actually used at this point (null for non-hyena). | +| `short_conv` | Hyena short-conv implementation actually used (null for non-hyena). | +| `mem_gb` | Peak CUDA memory allocated during the point. | + +## Reproducing + +```bash +# Reach sweeps (16 -> 16.7M): +for D in 1 2 3; do DATA_DIM=$D sbatch --time=04:30:00 --export=ALL \ + scripts/slurm/submit_forward_time_nd.sh; done + +# Flash-kernel sweeps: +for D in 1 2 3; do DATA_DIM=$D scripts/slurm/submit_forward_time_flash_kernels.sh; done +``` + +Both wrap `benchmarks/benchmark_forward_time_nd_resolution.py`. Relevant flags: +`--fft-backend {subq_ops,subq_ops_fused,torch_fft}`, `--short-conv {subq_ops,torch}`, +`--mixers`, `--resolutions`, `--hidden-dim`, `--num-iters`, `--max-seconds-per-point`. diff --git a/benchmarks/results/forward_time_1d.jsonl b/benchmarks/results/forward_time_1d.jsonl new file mode 100644 index 00000000..1fbc215e --- /dev/null +++ b/benchmarks/results/forward_time_1d.jsonl @@ -0,0 +1,63 @@ +{"mixer": "attention", "resolution": 16, "seq_len": 16, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.6519936243693034, "mem_gb": 0.03125905990600586, "iters": 30} +{"mixer": "attention", "resolution": 32, "seq_len": 32, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.1437323252360025, "mem_gb": 0.031262874603271484, "iters": 30} +{"mixer": "attention", "resolution": 64, "seq_len": 64, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.143288548787435, "mem_gb": 0.03127288818359375, "iters": 30} +{"mixer": "attention", "resolution": 128, "seq_len": 128, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.108474604288737, "mem_gb": 0.031293392181396484, "iters": 30} +{"mixer": "attention", "resolution": 256, "seq_len": 256, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.1490933736165365, "mem_gb": 0.03133440017700195, "iters": 30} +{"mixer": "attention", "resolution": 512, "seq_len": 512, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.1342474619547527, "mem_gb": 0.03166818618774414, "iters": 30} +{"mixer": "attention", "resolution": 1024, "seq_len": 1024, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.1777024586995444, "mem_gb": 0.032587528228759766, "iters": 30} +{"mixer": "attention", "resolution": 2048, "seq_len": 2048, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.6425365447998047, "mem_gb": 0.033922672271728516, "iters": 30} +{"mixer": "attention", "resolution": 4096, "seq_len": 4096, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.187606430053711, "mem_gb": 0.034578800201416016, "iters": 30} +{"mixer": "attention", "resolution": 8192, "seq_len": 8192, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.1390037536621094, "mem_gb": 0.033876895904541016, "iters": 30} +{"mixer": "attention", "resolution": 16384, "seq_len": 16384, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.092416000366211, "mem_gb": 0.036501407623291016, "iters": 30} +{"mixer": "attention", "resolution": 32768, "seq_len": 32768, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.394927978515625, "mem_gb": 0.041750431060791016, "iters": 30} +{"mixer": "attention", "resolution": 65536, "seq_len": 65536, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 5.02381337483724, "mem_gb": 0.052248477935791016, "iters": 30} +{"mixer": "attention", "resolution": 131072, "seq_len": 131072, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 17.16387939453125, "mem_gb": 0.07324457168579102, "iters": 30} +{"mixer": "attention", "resolution": 262144, "seq_len": 262144, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 67.81159261067708, "mem_gb": 0.11523675918579102, "iters": 30} +{"mixer": "attention", "resolution": 524288, "seq_len": 524288, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 261.96787623355266, "mem_gb": 0.19922113418579102, "iters": 19} +{"mixer": "attention", "resolution": 1048576, "seq_len": 1048576, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1046.6390625, "mem_gb": 0.367189884185791, "iters": 5} +{"mixer": "attention", "resolution": 2097152, "seq_len": 2097152, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 4185.4736328125, "mem_gb": 0.703127384185791, "iters": 3} +{"mixer": "attention", "resolution": 4194304, "seq_len": 4194304, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 16788.3125, "mem_gb": 1.375002384185791, "iters": 3} +{"mixer": "attention", "resolution": 8388608, "seq_len": 8388608, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 67089.92708333333, "mem_gb": 2.718752384185791, "iters": 3} +{"mixer": "attention", "resolution": 16777216, "seq_len": 16777216, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 268037.9375, "mem_gb": 5.406252384185791, "iters": 1} +{"mixer": "hyena", "resolution": 16, "seq_len": 16, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.6665813128153484, "mem_gb": 0.0312809944152832, "iters": 3} +{"mixer": "hyena", "resolution": 32, "seq_len": 32, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.5079882303873697, "mem_gb": 0.03128337860107422, "iters": 30} +{"mixer": "hyena", "resolution": 64, "seq_len": 64, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.5147754669189453, "mem_gb": 0.031290531158447266, "iters": 30} +{"mixer": "hyena", "resolution": 128, "seq_len": 128, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.4696756998697917, "mem_gb": 0.03131437301635742, "iters": 30} +{"mixer": "hyena", "resolution": 256, "seq_len": 256, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.4313332875569662, "mem_gb": 0.03136444091796875, "iters": 30} +{"mixer": "hyena", "resolution": 512, "seq_len": 512, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.4121524810791015, "mem_gb": 0.031464576721191406, "iters": 30} +{"mixer": "hyena", "resolution": 1024, "seq_len": 1024, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.4532885233561197, "mem_gb": 0.03166484832763672, "iters": 30} +{"mixer": "hyena", "resolution": 2048, "seq_len": 2048, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.4503093719482423, "mem_gb": 0.032065391540527344, "iters": 30} +{"mixer": "hyena", "resolution": 4096, "seq_len": 4096, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.4517504374186199, "mem_gb": 0.032866477966308594, "iters": 30} +{"mixer": "hyena", "resolution": 8192, "seq_len": 8192, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.4528917948404947, "mem_gb": 0.034468650817871094, "iters": 30} +{"mixer": "hyena", "resolution": 16384, "seq_len": 16384, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.4051092783610026, "mem_gb": 0.037672996520996094, "iters": 30} +{"mixer": "hyena", "resolution": 32768, "seq_len": 32768, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.4203936258951824, "mem_gb": 0.044081687927246094, "iters": 30} +{"mixer": "hyena", "resolution": 65536, "seq_len": 65536, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.4971509297688803, "mem_gb": 0.0637655258178711, "iters": 30} +{"mixer": "hyena", "resolution": 131072, "seq_len": 131072, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.5013355255126952, "mem_gb": 0.0923910140991211, "iters": 30} +{"mixer": "hyena", "resolution": 262144, "seq_len": 262144, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.4795188903808594, "mem_gb": 0.1534261703491211, "iters": 30} +{"mixer": "hyena", "resolution": 524288, "seq_len": 524288, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.9206677754720052, "mem_gb": 0.2755880355834961, "iters": 30} +{"mixer": "hyena", "resolution": 1048576, "seq_len": 1048576, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 4.34810536702474, "mem_gb": 0.5197286605834961, "iters": 30} +{"mixer": "hyena", "resolution": 2097152, "seq_len": 2097152, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 8.865488688151041, "mem_gb": 1.008193016052246, "iters": 30} +{"mixer": "hyena", "resolution": 4194304, "seq_len": 4194304, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 18.639249674479167, "mem_gb": 1.984755516052246, "iters": 30} +{"mixer": "hyena", "resolution": 8388608, "seq_len": 8388608, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 37.347111002604166, "mem_gb": 3.938246726989746, "iters": 30} +{"mixer": "hyena", "resolution": 16777216, "seq_len": 16777216, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 74.5071533203125, "mem_gb": 7.844496726989746, "iters": 30} +{"mixer": "mamba", "resolution": 16, "seq_len": 16, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.236959934234619, "mem_gb": 0.28156232833862305, "iters": 3} +{"mixer": "mamba", "resolution": 32, "seq_len": 32, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.4100128173828126, "mem_gb": 0.03158140182495117, "iters": 30} +{"mixer": "mamba", "resolution": 64, "seq_len": 64, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.407898712158203, "mem_gb": 0.03161811828613281, "iters": 30} +{"mixer": "mamba", "resolution": 128, "seq_len": 128, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.4106356302897134, "mem_gb": 0.031691551208496094, "iters": 30} +{"mixer": "mamba", "resolution": 256, "seq_len": 256, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.412281545003255, "mem_gb": 0.03183889389038086, "iters": 30} +{"mixer": "mamba", "resolution": 512, "seq_len": 512, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.271894327799479, "mem_gb": 0.03238534927368164, "iters": 30} +{"mixer": "mamba", "resolution": 1024, "seq_len": 1024, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.390685780843099, "mem_gb": 0.0334782600402832, "iters": 30} +{"mixer": "mamba", "resolution": 2048, "seq_len": 2048, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.4150848388671875, "mem_gb": 0.03566408157348633, "iters": 30} +{"mixer": "mamba", "resolution": 4096, "seq_len": 4096, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.2794708251953124, "mem_gb": 0.04003572463989258, "iters": 30} +{"mixer": "mamba", "resolution": 8192, "seq_len": 8192, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.425262959798177, "mem_gb": 0.04877901077270508, "iters": 30} +{"mixer": "mamba", "resolution": 16384, "seq_len": 16384, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.444763692220052, "mem_gb": 0.06626558303833008, "iters": 30} +{"mixer": "mamba", "resolution": 32768, "seq_len": 32768, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.438696543375651, "mem_gb": 0.10221529006958008, "iters": 30} +{"mixer": "mamba", "resolution": 65536, "seq_len": 65536, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.4371327718098956, "mem_gb": 0.17118501663208008, "iters": 30} +{"mixer": "mamba", "resolution": 131072, "seq_len": 131072, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.445636240641276, "mem_gb": 0.3110775947570801, "iters": 30} +{"mixer": "mamba", "resolution": 262144, "seq_len": 262144, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 3.0665514628092447, "mem_gb": 0.5918393135070801, "iters": 30} +{"mixer": "mamba", "resolution": 524288, "seq_len": 524288, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 5.956387329101562, "mem_gb": 1.15043306350708, "iters": 30} +{"mixer": "mamba", "resolution": 1048576, "seq_len": 1048576, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 11.771566772460938, "mem_gb": 2.26957368850708, "iters": 30} +{"mixer": "mamba", "resolution": 2097152, "seq_len": 2097152, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 23.35182088216146, "mem_gb": 4.50785493850708, "iters": 30} +{"mixer": "mamba", "resolution": 4194304, "seq_len": 4194304, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "error", "ms": null, "mem_gb": null, "iters": null, "detail": "AcceleratorError(\"CUDA error: invalid argument\\nSearch for `cudaErrorInvalidValue' 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": 8388608, "seq_len": 8388608, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "error", "ms": null, "mem_gb": null, "iters": null, "detail": "AcceleratorError(\"CUDA error: invalid argument\\nSearch for `cudaErrorInvalidValue' 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, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "error", "ms": null, "mem_gb": null, "iters": null, "detail": "AcceleratorError(\"CUDA error: invalid argument\\nSearch for `cudaErrorInvalidValue' 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/forward_time_2d.jsonl b/benchmarks/results/forward_time_2d.jsonl new file mode 100644 index 00000000..dcec862b --- /dev/null +++ b/benchmarks/results/forward_time_2d.jsonl @@ -0,0 +1,27 @@ +{"mixer": "attention", "resolution": 16, "seq_len": 256, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.8033141454060873, "mem_gb": 0.031328678131103516, "iters": 30} +{"mixer": "attention", "resolution": 32, "seq_len": 1024, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.4074293772379558, "mem_gb": 0.03255891799926758, "iters": 30} +{"mixer": "attention", "resolution": 64, "seq_len": 4096, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.4830282847086589, "mem_gb": 0.03445863723754883, "iters": 30} +{"mixer": "attention", "resolution": 128, "seq_len": 16384, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.4427925109863282, "mem_gb": 0.03601694107055664, "iters": 30} +{"mixer": "attention", "resolution": 256, "seq_len": 65536, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 5.0957794189453125, "mem_gb": 0.050302982330322266, "iters": 30} +{"mixer": "attention", "resolution": 512, "seq_len": 262144, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 67.89158935546875, "mem_gb": 0.10743951797485352, "iters": 30} +{"mixer": "attention", "resolution": 1024, "seq_len": 1048576, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1046.96328125, "mem_gb": 0.335970401763916, "iters": 5} +{"mixer": "attention", "resolution": 2048, "seq_len": 4194304, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 16787.928385416668, "mem_gb": 1.250063419342041, "iters": 3} +{"mixer": "attention", "resolution": 4096, "seq_len": 16777216, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 268220.8125, "mem_gb": 4.906374454498291, "iters": 1} +{"mixer": "hyena", "resolution": 16, "seq_len": 256, "data_dim": 2, "backend": "subq_ops_fused", "short_conv": "torch_conv2d_symmetric", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.2988800048828124, "mem_gb": 0.03142833709716797, "iters": 10} +{"mixer": "hyena", "resolution": 32, "seq_len": 1024, "data_dim": 2, "backend": "subq_ops_fused", "short_conv": "torch_conv2d_symmetric", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.113335715399848, "mem_gb": 0.03191518783569336, "iters": 27} +{"mixer": "hyena", "resolution": 64, "seq_len": 4096, "data_dim": 2, "backend": "subq_ops_fused", "short_conv": "torch_conv2d_symmetric", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.0932266800491899, "mem_gb": 0.03386116027832031, "iters": 27} +{"mixer": "hyena", "resolution": 128, "seq_len": 16384, "data_dim": 2, "backend": "subq_ops", "short_conv": "torch_conv2d_symmetric", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.1427618173452525, "mem_gb": 0.039937496185302734, "iters": 26} +{"mixer": "hyena", "resolution": 256, "seq_len": 65536, "data_dim": 2, "backend": "subq_ops", "short_conv": "torch_conv2d_symmetric", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.045998793381911, "mem_gb": 0.06593656539916992, "iters": 26} +{"mixer": "hyena", "resolution": 512, "seq_len": 262144, "data_dim": 2, "backend": "subq_ops", "short_conv": "torch_conv2d_symmetric", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.0553960005442302, "mem_gb": 0.1699366569519043, "iters": 24} +{"mixer": "hyena", "resolution": 1024, "seq_len": 1048576, "data_dim": 2, "backend": "subq_ops", "short_conv": "torch_conv2d_symmetric", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.360860263600069, "mem_gb": 0.585944652557373, "iters": 17} +{"mixer": "hyena", "resolution": 2048, "seq_len": 4194304, "data_dim": 2, "backend": "subq_ops", "short_conv": "torch_conv2d_symmetric", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 10.15641975402832, "mem_gb": 2.2500219345092773, "iters": 8} +{"mixer": "hyena", "resolution": 4096, "seq_len": 16777216, "data_dim": 2, "backend": "subq_ops", "short_conv": "torch_conv2d_symmetric", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 41.291552225748696, "mem_gb": 8.906271934509277, "iters": 3} +{"mixer": "mamba", "resolution": 16, "seq_len": 256, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.2885866165161133, "mem_gb": 0.281827449798584, "iters": 3} +{"mixer": "mamba", "resolution": 32, "seq_len": 1024, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.1538421630859377, "mem_gb": 0.0334782600402832, "iters": 30} +{"mixer": "mamba", "resolution": 64, "seq_len": 4096, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.1768618265787762, "mem_gb": 0.04003572463989258, "iters": 30} +{"mixer": "mamba", "resolution": 128, "seq_len": 16384, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.33219731648763, "mem_gb": 0.06626558303833008, "iters": 30} +{"mixer": "mamba", "resolution": 256, "seq_len": 65536, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.338206990559896, "mem_gb": 0.17118501663208008, "iters": 30} +{"mixer": "mamba", "resolution": 512, "seq_len": 262144, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 3.19793701171875, "mem_gb": 0.5918393135070801, "iters": 30} +{"mixer": "mamba", "resolution": 1024, "seq_len": 1048576, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 12.22380879720052, "mem_gb": 2.26957368850708, "iters": 30} +{"mixer": "mamba", "resolution": 2048, "seq_len": 4194304, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "error", "ms": null, "mem_gb": null, "iters": null, "detail": "AcceleratorError(\"CUDA error: invalid argument\\nSearch for `cudaErrorInvalidValue' 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, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "error", "ms": null, "mem_gb": null, "iters": null, "detail": "AcceleratorError(\"CUDA error: invalid argument\\nSearch for `cudaErrorInvalidValue' 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/forward_time_3d.jsonl b/benchmarks/results/forward_time_3d.jsonl new file mode 100644 index 00000000..593e72ff --- /dev/null +++ b/benchmarks/results/forward_time_3d.jsonl @@ -0,0 +1,15 @@ +{"mixer": "attention", "resolution": 16, "seq_len": 4096, "data_dim": 3, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.5790976206461589, "mem_gb": 0.034456729888916016, "iters": 30} +{"mixer": "attention", "resolution": 32, "seq_len": 32768, "data_dim": 3, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.334172821044922, "mem_gb": 0.040773868560791016, "iters": 30} +{"mixer": "attention", "resolution": 64, "seq_len": 262144, "data_dim": 3, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 67.67419840494792, "mem_gb": 0.10742425918579102, "iters": 30} +{"mixer": "attention", "resolution": 128, "seq_len": 2097152, "data_dim": 3, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 4178.6015625, "mem_gb": 0.640627384185791, "iters": 3} +{"mixer": "attention", "resolution": 256, "seq_len": 16777216, "data_dim": 3, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 267676.53125, "mem_gb": 4.906252384185791, "iters": 1} +{"mixer": "hyena", "resolution": 16, "seq_len": 4096, "data_dim": 3, "backend": "torch_fft", "short_conv": "torch_conv3d_symmetric", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.1183239618937175, "mem_gb": 0.033939361572265625, "iters": 24} +{"mixer": "hyena", "resolution": 32, "seq_len": 32768, "data_dim": 3, "backend": "torch_fft", "short_conv": "torch_conv3d_symmetric", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.1153109232584635, "mem_gb": 0.05238676071166992, "iters": 30} +{"mixer": "hyena", "resolution": 64, "seq_len": 262144, "data_dim": 3, "backend": "torch_fft", "short_conv": "torch_conv3d_symmetric", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.1217333475748699, "mem_gb": 0.20203018188476562, "iters": 30} +{"mixer": "hyena", "resolution": 128, "seq_len": 2097152, "data_dim": 3, "backend": "torch_fft", "short_conv": "torch_conv3d_symmetric", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 7.578884124755859, "mem_gb": 1.372342586517334, "iters": 16} +{"mixer": "hyena", "resolution": 256, "seq_len": 16777216, "data_dim": 3, "backend": "torch_fft", "short_conv": "torch_conv3d_symmetric", "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 62.03134536743164, "mem_gb": 10.737355709075928, "iters": 4} +{"mixer": "mamba", "resolution": 16, "seq_len": 4096, "data_dim": 3, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.4006826082865396, "mem_gb": 0.2898526191711426, "iters": 3} +{"mixer": "mamba", "resolution": 32, "seq_len": 32768, "data_dim": 3, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.413733418782552, "mem_gb": 0.10221529006958008, "iters": 30} +{"mixer": "mamba", "resolution": 64, "seq_len": 262144, "data_dim": 3, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 3.1418612162272135, "mem_gb": 0.5918393135070801, "iters": 30} +{"mixer": "mamba", "resolution": 128, "seq_len": 2097152, "data_dim": 3, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 23.84066162109375, "mem_gb": 4.50785493850708, "iters": 30} +{"mixer": "mamba", "resolution": 256, "seq_len": 16777216, "data_dim": 3, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 8, "num_heads": 2, "dtype": "bf16", "device": "NVIDIA GB200", "status": "error", "ms": null, "mem_gb": null, "iters": null, "detail": "AcceleratorError(\"CUDA error: invalid argument\\nSearch for `cudaErrorInvalidValue' 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/forward_time_flash_1d.jsonl b/benchmarks/results/forward_time_flash_1d.jsonl new file mode 100644 index 00000000..8c1fedb8 --- /dev/null +++ b/benchmarks/results/forward_time_flash_1d.jsonl @@ -0,0 +1,84 @@ +{"mixer": "hyena", "resolution": 16, "seq_len": 16, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.9558293024698894, "mem_gb": 0.03680849075317383, "iters": 3} +{"mixer": "hyena", "resolution": 32, "seq_len": 32, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.42827517191569, "mem_gb": 0.03690004348754883, "iters": 30} +{"mixer": "hyena", "resolution": 64, "seq_len": 64, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.3660415649414062, "mem_gb": 0.03708314895629883, "iters": 30} +{"mixer": "hyena", "resolution": 128, "seq_len": 128, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.3322218577067058, "mem_gb": 0.03757047653198242, "iters": 30} +{"mixer": "hyena", "resolution": 256, "seq_len": 256, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.3276842753092448, "mem_gb": 0.039402008056640625, "iters": 30} +{"mixer": "hyena", "resolution": 512, "seq_len": 512, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.3222677866617838, "mem_gb": 0.04355335235595703, "iters": 30} +{"mixer": "hyena", "resolution": 1024, "seq_len": 1024, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.3257119496663412, "mem_gb": 0.053320884704589844, "iters": 30} +{"mixer": "hyena", "resolution": 2048, "seq_len": 2048, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.2892011006673176, "mem_gb": 0.06846141815185547, "iters": 30} +{"mixer": "hyena", "resolution": 4096, "seq_len": 4096, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.2864607493082683, "mem_gb": 0.10264968872070312, "iters": 30} +{"mixer": "hyena", "resolution": 8192, "seq_len": 8192, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.29236691792806, "mem_gb": 0.16907119750976562, "iters": 30} +{"mixer": "hyena", "resolution": 16384, "seq_len": 16384, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.2895647684733074, "mem_gb": 0.3009376525878906, "iters": 30} +{"mixer": "hyena", "resolution": 32768, "seq_len": 32768, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.2863882700602214, "mem_gb": 0.5666236877441406, "iters": 30} +{"mixer": "hyena", "resolution": 65536, "seq_len": 65536, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.395428212483724, "mem_gb": 1.9788551330566406, "iters": 30} +{"mixer": "hyena", "resolution": 131072, "seq_len": 131072, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 4.761155192057291, "mem_gb": 3.9165992736816406, "iters": 30} +{"mixer": "hyena", "resolution": 262144, "seq_len": 262144, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 9.971072387695312, "mem_gb": 7.792087554931641, "iters": 30} +{"mixer": "hyena", "resolution": 524288, "seq_len": 524288, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 24.203129069010416, "mem_gb": 15.54892349243164, "iters": 30} +{"mixer": "hyena", "resolution": 1048576, "seq_len": 1048576, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 76.2148193359375, "mem_gb": 31.05087661743164, "iters": 30} +{"mixer": "attention", "resolution": 16, "seq_len": 16, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.5511717515833238, "mem_gb": 0.0367279052734375, "iters": 17} +{"mixer": "attention", "resolution": 32, "seq_len": 32, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.5583242115221525, "mem_gb": 0.036834716796875, "iters": 19} +{"mixer": "attention", "resolution": 64, "seq_len": 64, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.5617119789123535, "mem_gb": 0.03704833984375, "iters": 20} +{"mixer": "attention", "resolution": 128, "seq_len": 128, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.5478384017944335, "mem_gb": 0.0374755859375, "iters": 20} +{"mixer": "attention", "resolution": 256, "seq_len": 256, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.6171807765960693, "mem_gb": 0.038331031799316406, "iters": 10} +{"mixer": "attention", "resolution": 512, "seq_len": 512, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.6165791988372803, "mem_gb": 0.041504859924316406, "iters": 10} +{"mixer": "attention", "resolution": 1024, "seq_len": 1024, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.6142367839813232, "mem_gb": 0.047852516174316406, "iters": 10} +{"mixer": "attention", "resolution": 2048, "seq_len": 2048, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.6111775875091553, "mem_gb": 0.060547828674316406, "iters": 10} +{"mixer": "attention", "resolution": 4096, "seq_len": 4096, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.6322559780544705, "mem_gb": 0.0869150161743164, "iters": 9} +{"mixer": "attention", "resolution": 8192, "seq_len": 8192, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.6307733323838975, "mem_gb": 0.1376962661743164, "iters": 9} +{"mixer": "attention", "resolution": 16384, "seq_len": 16384, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.8223928875393338, "mem_gb": 0.2392587661743164, "iters": 9} +{"mixer": "attention", "resolution": 32768, "seq_len": 32768, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.22761959499783, "mem_gb": 0.4414072036743164, "iters": 9} +{"mixer": "attention", "resolution": 65536, "seq_len": 65536, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 6.9648924933539496, "mem_gb": 0.8476572036743164, "iters": 9} +{"mixer": "attention", "resolution": 131072, "seq_len": 131072, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 24.98052300347222, "mem_gb": 1.6601572036743164, "iters": 9} +{"mixer": "attention", "resolution": 262144, "seq_len": 262144, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 95.3530502319336, "mem_gb": 3.2851572036743164, "iters": 8} +{"mixer": "attention", "resolution": 524288, "seq_len": 524288, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 379.8583984375, "mem_gb": 6.535157203674316, "iters": 5} +{"mixer": "attention", "resolution": 1048576, "seq_len": 1048576, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1587.4446614583333, "mem_gb": 13.035157203674316, "iters": 3} +{"mixer": "flex", "resolution": 16, "seq_len": 16, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.9973440170288086, "mem_gb": 0.037711143493652344, "iters": 3} +{"mixer": "flex", "resolution": 32, "seq_len": 32, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.8933866818745931, "mem_gb": 0.16705989837646484, "iters": 3} +{"mixer": "flex", "resolution": 64, "seq_len": 64, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.0686346689860027, "mem_gb": 0.04506874084472656, "iters": 30} +{"mixer": "flex", "resolution": 128, "seq_len": 128, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.9812880158424377, "mem_gb": 0.0380859375, "iters": 4} +{"mixer": "flex", "resolution": 256, "seq_len": 256, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.0236010869344077, "mem_gb": 0.038330078125, "iters": 30} +{"mixer": "flex", "resolution": 512, "seq_len": 512, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.0485482533772787, "mem_gb": 0.04052734375, "iters": 30} +{"mixer": "flex", "resolution": 1024, "seq_len": 1024, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.0429706573486328, "mem_gb": 0.0458984375, "iters": 30} +{"mixer": "flex", "resolution": 2048, "seq_len": 2048, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.0357813517252603, "mem_gb": 0.056640625, "iters": 30} +{"mixer": "flex", "resolution": 4096, "seq_len": 4096, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.04617493947347, "mem_gb": 0.078125, "iters": 30} +{"mixer": "flex", "resolution": 8192, "seq_len": 8192, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.0557600021362306, "mem_gb": 0.1220703125, "iters": 30} +{"mixer": "flex", "resolution": 16384, "seq_len": 16384, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.424748738606771, "mem_gb": 0.2080078125, "iters": 30} +{"mixer": "flex", "resolution": 32768, "seq_len": 32768, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 7.597084554036458, "mem_gb": 0.37890625, "iters": 30} +{"mixer": "flex", "resolution": 65536, "seq_len": 65536, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 28.295279947916665, "mem_gb": 0.72265625, "iters": 30} +{"mixer": "flex", "resolution": 131072, "seq_len": 131072, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 114.48787434895833, "mem_gb": 1.41015625, "iters": 30} +{"mixer": "flex", "resolution": 262144, "seq_len": 262144, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 460.90860262784093, "mem_gb": 2.78515625, "iters": 11} +{"mixer": "flex", "resolution": 524288, "seq_len": 524288, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1831.2366536458333, "mem_gb": 5.53515625, "iters": 3} +{"mixer": "flex", "resolution": 1048576, "seq_len": 1048576, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 7299.468098958333, "mem_gb": 11.03515625, "iters": 3} +{"mixer": "fa4", "resolution": 16, "seq_len": 16, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.8252693017323812, "mem_gb": 0.0367279052734375, "iters": 3} +{"mixer": "fa4", "resolution": 32, "seq_len": 32, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.9502549489339193, "mem_gb": 0.036834716796875, "iters": 30} +{"mixer": "fa4", "resolution": 64, "seq_len": 64, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.948692258199056, "mem_gb": 0.03704833984375, "iters": 30} +{"mixer": "fa4", "resolution": 128, "seq_len": 128, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.9553461074829102, "mem_gb": 0.0374755859375, "iters": 30} +{"mixer": "fa4", "resolution": 256, "seq_len": 256, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.7585386435190836, "mem_gb": 0.038330078125, "iters": 3} +{"mixer": "fa4", "resolution": 512, "seq_len": 512, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.7824107011159261, "mem_gb": 0.04052734375, "iters": 3} +{"mixer": "fa4", "resolution": 1024, "seq_len": 1024, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.9382858912150065, "mem_gb": 0.0458984375, "iters": 30} +{"mixer": "fa4", "resolution": 2048, "seq_len": 2048, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.9607999801635743, "mem_gb": 0.056640625, "iters": 30} +{"mixer": "fa4", "resolution": 4096, "seq_len": 4096, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.9603445053100585, "mem_gb": 0.078125, "iters": 30} +{"mixer": "fa4", "resolution": 8192, "seq_len": 8192, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.9593941370646158, "mem_gb": 0.1220703125, "iters": 30} +{"mixer": "fa4", "resolution": 16384, "seq_len": 16384, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.9963861465454101, "mem_gb": 0.2080078125, "iters": 30} +{"mixer": "fa4", "resolution": 32768, "seq_len": 32768, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.6575691223144533, "mem_gb": 0.37890625, "iters": 30} +{"mixer": "fa4", "resolution": 65536, "seq_len": 65536, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 8.120289103190105, "mem_gb": 0.72265625, "iters": 30} +{"mixer": "fa4", "resolution": 131072, "seq_len": 131072, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 30.295428466796874, "mem_gb": 1.41015625, "iters": 30} +{"mixer": "fa4", "resolution": 262144, "seq_len": 262144, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 119.90684407552084, "mem_gb": 2.78515625, "iters": 30} +{"mixer": "fa4", "resolution": 524288, "seq_len": 524288, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 505.537158203125, "mem_gb": 5.53515625, "iters": 10} +{"mixer": "fa4", "resolution": 1048576, "seq_len": 1048576, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2085.9620768229165, "mem_gb": 11.03515625, "iters": 3} +{"mixer": "mamba", "resolution": 16, "seq_len": 16, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.2863146464029946, "mem_gb": 0.2954540252685547, "iters": 3} +{"mixer": "mamba", "resolution": 32, "seq_len": 32, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.2907915751139325, "mem_gb": 0.04659843444824219, "iters": 30} +{"mixer": "mamba", "resolution": 64, "seq_len": 64, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.298900349934896, "mem_gb": 0.046919822692871094, "iters": 30} +{"mixer": "mamba", "resolution": 128, "seq_len": 128, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.3209856669108073, "mem_gb": 0.04788398742675781, "iters": 30} +{"mixer": "mamba", "resolution": 256, "seq_len": 256, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.296706136067708, "mem_gb": 0.04996347427368164, "iters": 30} +{"mixer": "mamba", "resolution": 512, "seq_len": 512, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.1865216573079427, "mem_gb": 0.0543971061706543, "iters": 30} +{"mixer": "mamba", "resolution": 1024, "seq_len": 1024, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.3311030069986978, "mem_gb": 0.06369161605834961, "iters": 30} +{"mixer": "mamba", "resolution": 2048, "seq_len": 2048, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.358390299479167, "mem_gb": 0.08282995223999023, "iters": 30} +{"mixer": "mamba", "resolution": 4096, "seq_len": 4096, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.189985148111979, "mem_gb": 0.11646795272827148, "iters": 30} +{"mixer": "mamba", "resolution": 8192, "seq_len": 8192, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.2974624633789062, "mem_gb": 0.18740606307983398, "iters": 30} +{"mixer": "mamba", "resolution": 16384, "seq_len": 16384, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.321709950764974, "mem_gb": 0.329282283782959, "iters": 30} +{"mixer": "mamba", "resolution": 32768, "seq_len": 32768, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.3100575764973956, "mem_gb": 0.614011287689209, "iters": 30} +{"mixer": "mamba", "resolution": 65536, "seq_len": 65536, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 3.0068618774414064, "mem_gb": 1.180539608001709, "iters": 30} +{"mixer": "mamba", "resolution": 131072, "seq_len": 131072, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 5.8191680908203125, "mem_gb": 2.315549373626709, "iters": 30} +{"mixer": "mamba", "resolution": 262144, "seq_len": 262144, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 11.45050048828125, "mem_gb": 4.585568904876709, "iters": 30} +{"mixer": "mamba", "resolution": 524288, "seq_len": 524288, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 22.791756184895835, "mem_gb": 9.125607967376709, "iters": 30} diff --git a/benchmarks/results/forward_time_flash_2d.jsonl b/benchmarks/results/forward_time_flash_2d.jsonl new file mode 100644 index 00000000..ed406196 --- /dev/null +++ b/benchmarks/results/forward_time_flash_2d.jsonl @@ -0,0 +1,34 @@ +{"mixer": "hyena", "resolution": 16, "seq_len": 256, "data_dim": 2, "backend": "subq_ops_fused", "short_conv": "torch_conv2d_symmetric", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.6246639490127563, "mem_gb": 0.04557514190673828, "iters": 8} +{"mixer": "hyena", "resolution": 32, "seq_len": 1024, "data_dim": 2, "backend": "subq_ops_fused", "short_conv": "torch_conv2d_symmetric", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.2256444295247395, "mem_gb": 0.0779581069946289, "iters": 18} +{"mixer": "hyena", "resolution": 64, "seq_len": 4096, "data_dim": 2, "backend": "subq_ops_fused", "short_conv": "torch_conv2d_symmetric", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.2114700078964233, "mem_gb": 0.20084333419799805, "iters": 16} +{"mixer": "hyena", "resolution": 128, "seq_len": 16384, "data_dim": 2, "backend": "subq_ops", "short_conv": "torch_conv2d_symmetric", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.3589375813802083, "mem_gb": 0.5827198028564453, "iters": 15} +{"mixer": "hyena", "resolution": 256, "seq_len": 65536, "data_dim": 2, "backend": "subq_ops", "short_conv": "torch_conv2d_symmetric", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 3.5547008514404297, "mem_gb": 2.2237091064453125, "iters": 10} +{"mixer": "hyena", "resolution": 512, "seq_len": 262144, "data_dim": 2, "backend": "subq_ops", "short_conv": "torch_conv2d_symmetric", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 14.074893188476562, "mem_gb": 8.787670135498047, "iters": 5} +{"mixer": "hyena", "resolution": 1024, "seq_len": 1048576, "data_dim": 2, "backend": "subq_ops", "short_conv": "torch_conv2d_symmetric", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 89.18731689453125, "mem_gb": 29.543086528778076, "iters": 3} +{"mixer": "attention", "resolution": 16, "seq_len": 256, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.9855075412326388, "mem_gb": 0.038102149963378906, "iters": 9} +{"mixer": "attention", "resolution": 32, "seq_len": 1024, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.9488448143005371, "mem_gb": 0.046906471252441406, "iters": 10} +{"mixer": "attention", "resolution": 64, "seq_len": 4096, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.9660160276624892, "mem_gb": 0.0830698013305664, "iters": 9} +{"mixer": "attention", "resolution": 128, "seq_len": 16384, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.9985564549763998, "mem_gb": 0.2237558364868164, "iters": 9} +{"mixer": "attention", "resolution": 256, "seq_len": 65536, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 7.263189527723524, "mem_gb": 0.7854013442993164, "iters": 9} +{"mixer": "attention", "resolution": 512, "seq_len": 262144, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 96.78620147705078, "mem_gb": 3.0356454849243164, "iters": 8} +{"mixer": "attention", "resolution": 1024, "seq_len": 1048576, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1648.5638020833333, "mem_gb": 12.036133766174316, "iters": 3} +{"mixer": "flex", "resolution": 16, "seq_len": 256, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.2417279879252117, "mem_gb": 0.0393218994140625, "iters": 3} +{"mixer": "flex", "resolution": 32, "seq_len": 1024, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.3911040623982747, "mem_gb": 0.0508270263671875, "iters": 3} +{"mixer": "flex", "resolution": 64, "seq_len": 4096, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.5435595194498697, "mem_gb": 0.0792236328125, "iters": 30} +{"mixer": "flex", "resolution": 128, "seq_len": 16384, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.609130605061849, "mem_gb": 0.2083740234375, "iters": 30} +{"mixer": "flex", "resolution": 256, "seq_len": 65536, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 28.950142415364585, "mem_gb": 0.723876953125, "iters": 30} +{"mixer": "flex", "resolution": 512, "seq_len": 262144, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 467.33247514204544, "mem_gb": 2.78955078125, "iters": 11} +{"mixer": "flex", "resolution": 1024, "seq_len": 1048576, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 7381.654296875, "mem_gb": 11.0517578125, "iters": 3} +{"mixer": "fa4", "resolution": 16, "seq_len": 256, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.065130631128947, "mem_gb": 0.0381011962890625, "iters": 3} +{"mixer": "fa4", "resolution": 32, "seq_len": 1024, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.0374933083852131, "mem_gb": 0.0459442138671875, "iters": 3} +{"mixer": "fa4", "resolution": 64, "seq_len": 4096, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.4482111612955728, "mem_gb": 0.0792236328125, "iters": 30} +{"mixer": "fa4", "resolution": 128, "seq_len": 16384, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.4581642150878906, "mem_gb": 0.2083740234375, "iters": 30} +{"mixer": "fa4", "resolution": 256, "seq_len": 65536, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 8.411671447753907, "mem_gb": 0.723876953125, "iters": 30} +{"mixer": "fa4", "resolution": 512, "seq_len": 262144, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 122.08134765625, "mem_gb": 2.78955078125, "iters": 30} +{"mixer": "fa4", "resolution": 1024, "seq_len": 1048576, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2080.8497721354165, "mem_gb": 11.0517578125, "iters": 3} +{"mixer": "mamba", "resolution": 16, "seq_len": 256, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.5622080167134604, "mem_gb": 0.29776525497436523, "iters": 3} +{"mixer": "mamba", "resolution": 32, "seq_len": 1024, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.5105387369791665, "mem_gb": 0.06369161605834961, "iters": 30} +{"mixer": "mamba", "resolution": 64, "seq_len": 4096, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.5508852640787762, "mem_gb": 0.11646795272827148, "iters": 30} +{"mixer": "mamba", "resolution": 128, "seq_len": 16384, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.7360267639160156, "mem_gb": 0.329282283782959, "iters": 30} +{"mixer": "mamba", "resolution": 256, "seq_len": 65536, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.725147755940755, "mem_gb": 1.180539608001709, "iters": 30} +{"mixer": "mamba", "resolution": 512, "seq_len": 262144, "data_dim": 2, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 8.658978271484376, "mem_gb": 4.585568904876709, "iters": 30} diff --git a/benchmarks/results/forward_time_flash_3d.jsonl b/benchmarks/results/forward_time_flash_3d.jsonl new file mode 100644 index 00000000..581e894b --- /dev/null +++ b/benchmarks/results/forward_time_flash_3d.jsonl @@ -0,0 +1,15 @@ +{"mixer": "hyena", "resolution": 16, "seq_len": 4096, "data_dim": 3, "backend": "torch_fft", "short_conv": "torch_conv3d_symmetric", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.135732650756836, "mem_gb": 0.20885658264160156, "iters": 17} +{"mixer": "hyena", "resolution": 32, "seq_len": 32768, "data_dim": 3, "backend": "torch_fft", "short_conv": "torch_conv3d_symmetric", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 7.896506861636513, "mem_gb": 1.369001865386963, "iters": 19} +{"mixer": "hyena", "resolution": 64, "seq_len": 262144, "data_dim": 3, "backend": "torch_fft", "short_conv": "torch_conv3d_symmetric", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 70.85182189941406, "mem_gb": 10.609042644500732, "iters": 4} +{"mixer": "attention", "resolution": 16, "seq_len": 4096, "data_dim": 3, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.4962204297383626, "mem_gb": 0.0820322036743164, "iters": 9} +{"mixer": "attention", "resolution": 32, "seq_len": 32768, "data_dim": 3, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.7641208436754015, "mem_gb": 0.4101572036743164, "iters": 9} +{"mixer": "attention", "resolution": 64, "seq_len": 262144, "data_dim": 3, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 96.80563354492188, "mem_gb": 3.0351572036743164, "iters": 8} +{"mixer": "flex", "resolution": 16, "seq_len": 4096, "data_dim": 3, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.730026642481486, "mem_gb": 0.09375, "iters": 3} +{"mixer": "flex", "resolution": 32, "seq_len": 32768, "data_dim": 3, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 7.329098383585612, "mem_gb": 0.50390625, "iters": 3} +{"mixer": "flex", "resolution": 64, "seq_len": 262144, "data_dim": 3, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 459.91432883522725, "mem_gb": 2.53515625, "iters": 11} +{"mixer": "fa4", "resolution": 16, "seq_len": 4096, "data_dim": 3, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.708458662033081, "mem_gb": 0.07421875, "iters": 3} +{"mixer": "fa4", "resolution": 32, "seq_len": 32768, "data_dim": 3, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.150055440266927, "mem_gb": 0.34765625, "iters": 30} +{"mixer": "fa4", "resolution": 64, "seq_len": 262144, "data_dim": 3, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 118.7666015625, "mem_gb": 2.53515625, "iters": 30} +{"mixer": "mamba", "resolution": 16, "seq_len": 4096, "data_dim": 3, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.7073065439860025, "mem_gb": 0.345944881439209, "iters": 3} +{"mixer": "mamba", "resolution": 32, "seq_len": 32768, "data_dim": 3, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.9285003662109377, "mem_gb": 0.614011287689209, "iters": 30} +{"mixer": "mamba", "resolution": 64, "seq_len": 262144, "data_dim": 3, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 8.324334716796875, "mem_gb": 4.585568904876709, "iters": 30} diff --git a/benchmarks/results/nemotron_1d_e236.jsonl b/benchmarks/results/nemotron_1d_e236.jsonl new file mode 100644 index 00000000..59eb35c8 --- /dev/null +++ b/benchmarks/results/nemotron_1d_e236.jsonl @@ -0,0 +1,4 @@ +{"mixer": "hyena", "resolution": 8192, "seq_len": 8192, "data_dim": 1, "backend": "subq_ops", "short_conv": "subq_ops_causal_conv1d", "batch_size": 1, "hidden_dim": 768, "num_heads": 12, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.2762346903483073, "mem_gb": 0.5023279190063477, "iters": 30} +{"mixer": "mamba_causal", "resolution": 8192, "seq_len": 8192, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 768, "num_heads": 12, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.9675872166951497, "mem_gb": 0.5573253631591797, "iters": 30} +{"mixer": "mamba", "resolution": 8192, "seq_len": 8192, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 768, "num_heads": 12, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.0477301279703775, "mem_gb": 0.37429380416870117, "iters": 30} +{"mixer": "gdp", "resolution": 8192, "seq_len": 8192, "data_dim": 1, "backend": null, "short_conv": null, "batch_size": 1, "hidden_dim": 768, "num_heads": 12, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.292195129394531, "mem_gb": 1.528120994567871, "iters": 30} diff --git a/docs/benchmarks.md b/docs/benchmarks.md index ecf26d66..2892b902 100644 --- a/docs/benchmarks.md +++ b/docs/benchmarks.md @@ -38,7 +38,18 @@ official `mamba_chunk_scan_combined` Mamba2 kernel, and `nSubQ`: ![Forward time vs. sequence length for HyenaND, Attention, and Mamba](_static/throughput_scaling.png) HyenaND reaches million-token sequences at 265 ms (1M tokens), while attention -takes ~90 s (a **339×** gap) and the Mamba2 kernel runs out of memory. +takes ~90 s (a **339×** gap) and the Mamba2 kernel fails before reaching 1M. + +```{note} +**The Mamba2 curve stops for an implementation reason, not a memory one.** An +earlier revision of this page said the kernel "runs out of memory"; measured on a +GB200 (184 GiB) with mamba-ssm 2.3.2.post1, it does not. It fails with a CUDA +*illegal memory access* while using 22 GiB — 12% of available memory — because the +Triton kernels index with 32-bit arithmetic and overflow at ~2^31 elements. See +[Mamba-2 sequence-length limits](mamba2_limits.md) for the diagnosis and the +reproduction. The observed reach still matches this plot; only the stated cause +was wrong. +``` ## CUDA kernels (`nSubQ`) diff --git a/docs/index.rst b/docs/index.rst index 8ed06766..bb6e6ae3 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -148,6 +148,9 @@ Where to go next the base-config + ablation workflow. - :doc:`Benchmarks `: FLOP scaling, kernel speedups, and a worked ViT-5-Small ImageNet training optimization case study. +- :doc:`Mamba-2 Sequence-Length Limits `: why the Mamba-2 + baseline stops early in the forward-time sweeps — two measured implementation + limits, neither of them memory. - :doc:`Reports `: long-form technical reports backed by reproducible scripts and figures. - :doc:`Glossary `: quick definitions for SIREN, FiLM, implicit @@ -192,6 +195,7 @@ CUDA kernels live in a separate library: Lazy-Config System Examples Benchmarks + Mamba-2 Sequence-Length Limits Reports Glossary API Reference diff --git a/docs/mamba2_limits.md b/docs/mamba2_limits.md new file mode 100644 index 00000000..47a7713b --- /dev/null +++ b/docs/mamba2_limits.md @@ -0,0 +1,179 @@ +# Mamba-2 sequence-length limits + +Why the Mamba-2 baseline stops early in our forward-time sweeps, measured rather +than assumed. Two **separate** implementation limits, neither of them memory, and +neither inherent to the SSD algorithm. + +This page exists because [`benchmarks.md`](benchmarks.md) previously attributed the +stop to the kernel running out of memory. It does not run out of memory. The +*reach* reported there is reproducible; the stated cause was wrong. + +Measured on **GB200 (184 GiB)** with **mamba-ssm 2.3.2.post1**, causal-conv1d +1.6.2.post1, triton 3.7.1, torch 2.12.1+cu130 — 2.3.2.post1 being the latest +release on PyPI at the time of writing. + +______________________________________________________________________ + +## Summary + +| Limit | Where | Bound | Scales with | +| ----------------------- | -------------------------------------------------- | ---------------------------------- | --------------------------------------------------------- | +| 32-bit index overflow | `ssd_chunk_state.py` → `_chunk_state_fwd` (Triton) | ~231 **elements** | model width — longer sequences allowed at narrower widths | +| CUDA grid-dimension cap | `causal_conv1d` channels-last path | **4,194,240 tokens** = 65,535 × 64 | nothing — a fixed token count | + +**No configuration we tested ran out of memory.** Peak allocation never exceeded +87 GiB of 184 GiB, and at the 1M failure it was 22 GiB (12%). + +______________________________________________________________________ + +## Limit 1 — 32-bit index overflow in the SSD scan + +At realistic widths this is what stops Mamba-2 first. The fault site, with +`CUDA_LAUNCH_BLOCKING=1`: + +``` +mamba2.forward → mamba_split_conv1d_scan_combined + → _mamba_chunk_scan_combined_fwd (ssd_combined.py:375) + → _chunk_state_fwd (ssd_chunk_state.py:830) + → triton autotuner → do_bench → kernel_call +RuntimeError: Triton Error [CUDA]: an illegal memory access was encountered +``` + +An *illegal memory access* is a memory-safety fault — a thread dereferenced an +out-of-bounds address. It is not an allocator failure, and it is not the same +class of event as an OOM. + +### Evidence: the threshold scales inversely with width + +Bisected at `headdim=64, expand=2, d_state=128, ngroups=8`: + +| hidden | widest tensor width | last OK L | first FAIL L | elements at last OK | +| -----: | ------------------: | --------: | -----------: | ------------------------------------: | +| 768 | 3,584 | 598,016 | 606,208 | 2.143e9 = **99.8% of 231** | +| 1536 | 8,240 | 253,952 | 262,144 | 2.093e9 = **97.5% of 231** | + +Halving the reach when width doubles is the signature of an **element-count** +limit. A fixed sequence-length cap would not move with width; a memory limit +would show as an OOM at a much larger footprint. Both boundaries land just under +231 elements of the widest tensor in play. + +(The *binding* tensor differs between the two configs, so the claim here is +"~231 elements", not a pinpointed expression.) + +### Confirmed by patching + +Casting `tl.program_id(...)` to `tl.int64` before it feeds pointer arithmetic +fixes the 1M case: + +| Configuration | L = 1,048,576 | +| ----------------------------------------------------- | ------------------------------------- | +| Stock 2.3.2.post1 | illegal memory access | +| + int64 in `ssd_chunk_scan.py` + `layernorm_gated.py` | illegal memory access (**no change**) | +| + int64 also in `ssd_chunk_state.py` | **passes** | + +Necessary but **not sufficient**: with all three files patched, L=2,097,152 still +fails, with a Triton `invalid argument` — a launch-configuration rejection (grid +geometry), a different limit behind this one. + +### Upstream status + +Reported as [state-spaces/mamba#686](https://github.com/state-spaces/mamba/issues/686) +("Long Sequence Length Inference Mamba2: CUDA error: an illegal memory access was +encountered"), still open. As of this writing there is no fix: + +- 2.3.2.post1 is the latest PyPI release. +- `ssd_chunk_state.py` and `ssd_chunk_scan.py` are byte-identical between the + `v2.3.2` tag and `main`. +- The `mamba_ssm` tree contains no int64 index widening at all — the only two + `int64` occurrences are a dropout-seed dtype in `layer_norm.py` and a comment in + `mamba3_mimo.py`. + +The `wdykas@bfec072` ("more int64") commit referenced in that thread **does not +fix this crash**: it widens `program_id` in `ssd_chunk_scan.py` and +`layernorm_gated.py`, not in `ssd_chunk_state.py` where the fault occurs. We +verified this by applying a strict superset of its changes to both files (40 casts +versus its 4) — L=1M still failed. + +______________________________________________________________________ + +## Limit 2 — CUDA grid-dimension cap in `causal_conv1d` + +At small widths the SSD scan survives much further, and a different limit binds: +the short conv. It appears as `CUDA error: invalid argument` +(`cudaErrorInvalidValue`) — the driver *rejecting the launch*, so nothing executes. + +The boundary is exact and independent of channel count: + +| L | blocks (L / 64) | result | +| ------------: | --------------: | ------ | +| 4,194,176 | 65,534 | ok | +| **4,194,240** | **65,535** | **ok** | +| 4,194,304 | 65,536 | fail | + +`65,535` is the CUDA cap on a grid dimension, and the channels-last kernel tiles +the sequence at 64 elements per block, giving **L ≤ 65,535 × 64 = 4,194,240**. + +It only affects the **channels-last** path — which is the layout Mamba-2 passes, +via `rearrange(ensure_stride(xBC), "b s d -> b d s")`. A contiguous `[b, d, s]` +tensor runs to at least 32M tokens at every width we tried. + +______________________________________________________________________ + +## Why this is not an architectural limit + +The SSD chunk-scan is O(L) with fixed state per chunk; nothing in the algorithm +requires 32-bit offsets or caps a grid axis. Both limits are properties of the +current kernels: + +- the overflow is a pointer-arithmetic width choice, liftable with int64 casts + (demonstrated above); +- the grid cap is an unhandled input size in one layout, and the other layout has + no such bound. + +On the same GPU at the same widths, HyenaND reaches **16.7M tokens**. + +So these results say *"the shipping Mamba-2 kernels do not reach these lengths +today"* — a statement about an implementation, with a known and unmerged fix. They +do **not** support the stronger claim that Mamba-2 as an architecture cannot reach +them, and should not be cited that way. + +______________________________________________________________________ + +## Reproducing + +Both limits reproduce in a few lines. Mamba-2's package `__init__` eagerly imports +`mamba3` → `tilelang`, which collides with the `apache-tvm-ffi` that +subquadratic-ops-torch >= 0.2.2 installs, hence the shim. + +```python +import sys, torch + +sys.modules.setdefault("tilelang", None) # see note above + +# Limit 1 — overflow in the SSD scan (fails ~600K at hidden 768) +from mamba_ssm import Mamba2 + +m = Mamba2( + d_model=768, + headdim=64, + expand=2, + d_state=128, + ngroups=8, + device="cuda", + dtype=torch.bfloat16, +) +x = torch.randn(1, 1_048_576, 768, device="cuda", dtype=torch.bfloat16) +with torch.inference_mode(): + m(x) # illegal memory access, ~22 GiB peak + +# Limit 2 — grid cap in causal_conv1d, channels-last only +from causal_conv1d import causal_conv1d_fn + +w = torch.randn(272, 4, device="cuda", dtype=torch.bfloat16) +for L in (4_194_240, 4_194_304): # 65,535 vs 65,536 blocks + xc = torch.randn(1, L, 272, device="cuda", dtype=torch.bfloat16).transpose(1, 2) + causal_conv1d_fn(x=xc, weight=w, bias=None, activation="silu") +``` + +The sweeps that surfaced this are `scripts/slurm/submit_forward_time_nd.sh` (the +`mamba` / `mamba_causal` series) and `scripts/slurm/submit_nemotron_1d.sh`. diff --git a/nvsubquadratic/modules/attention.py b/nvsubquadratic/modules/attention.py index addc59f3..a590a422 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,39 @@ 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 + # Validate the unsupported argument combinations before touching the + # back-end imports below: `fa4` resolves an optional dependency that may + # raise ImportError, which would otherwise mask the clearer message here. + if attn_impl == "flex" and is_causal: + # forward() calls flex_attention without a BlockMask, so it cannot mask. + # Rejecting beats silently handing back bidirectional attention to a + # caller who asked for causal. + raise ValueError( + "attn_impl='flex' does not support is_causal=True: no causal BlockMask is built. " + "Use attn_impl='sdpa' or 'fa4' for causal attention." + ) + if attn_impl in ("flex", "fa4") and attn_dropout > 0.0: + # Neither flash-class path forwards a dropout probability — FlexAttention + # is called without one and FA3+ dropped the argument — so a non-zero + # value would be silently ignored rather than applied. + raise ValueError( + f"attn_impl={attn_impl!r} does not support attn_dropout={attn_dropout} (>0); " + "it would be silently ignored. Use attn_impl='sdpa' for attention dropout." + ) + + 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 +617,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/nvsubquadratic/modules/sequence_mixer.py b/nvsubquadratic/modules/sequence_mixer.py index 9ad59fd0..73ac2365 100644 --- a/nvsubquadratic/modules/sequence_mixer.py +++ b/nvsubquadratic/modules/sequence_mixer.py @@ -134,12 +134,22 @@ def __init__( out_proj_bias: bool = False, init_method_in: Callable[[int], Callable[[torch.Tensor], torch.Tensor]] | None = None, init_method_out: Callable[[int], Callable[[torch.Tensor], torch.Tensor]] | None = None, + inner_dim: int | None = None, ): """Initialise the QKV sequence mixer. Args: hidden_dim: Channel dimension ``C`` of the input / output tensor. - Both ``qkv_proj`` and ``out_proj`` are sized using this value. + ``qkv_proj`` reads and ``out_proj`` writes this width. + inner_dim: Width the inner mixer operates at, ``C_inner``. Defaults to + ``hidden_dim`` (the previous fixed behaviour). Setting it wider gives + the expanded parameterisation used for FLOP-matching against Mamba-2: + ``qkv_proj: C -> 3*C_inner``, mixer on ``C_inner`` channels, + ``out_proj: C_inner -> C``. With ``C_inner = e*C`` this is the ``e`` + of the FLOP-matching analysis; ``e = 1`` reproduces evo2's default, + which is only ~0.42x a Mamba-2 layer. **The inner mixer must be + constructed for ``inner_dim`` channels** — this class sizes the + projections, not the mixer. mixer_cfg: :class:`~nvsubquadratic.lazy_config.LazyConfig` for the inner sequence-mixing operator. The target class's ``forward`` method must accept ``(q, k, v, cp_group, **kwargs)`` where @@ -187,8 +197,11 @@ class cannot be constructed (e.g. missing required arguments or self.mixer = instantiate(mixer_cfg) - self.qkv_proj = torch.nn.Linear(hidden_dim, 3 * hidden_dim, bias=qkv_bias) - self.out_proj = torch.nn.Linear(hidden_dim, hidden_dim, bias=out_proj_bias) + self.hidden_dim = hidden_dim + self.inner_dim = hidden_dim if inner_dim is None else inner_dim + + self.qkv_proj = torch.nn.Linear(hidden_dim, 3 * self.inner_dim, bias=qkv_bias) + self.out_proj = torch.nn.Linear(self.inner_dim, hidden_dim, bias=out_proj_bias) if init_method_in is not None: init_method_in(hidden_dim)(self.qkv_proj.weight.data) 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..76572d3a --- /dev/null +++ b/scripts/slurm/submit_forward_time_flash_kernels.sh @@ -0,0 +1,68 @@ +#!/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}" + +# Powers of two from a 16-wide grid up to the largest R keeping 3*HIDDEN_DIM*R^DATA_DIM +# under torch's 2^31 index limit at hidden 512 (~1.4M tokens). +case "${DATA_DIM}" in + 1) RESOLUTIONS="${RESOLUTIONS:-16 32 64 128 256 512 1024 2048 4096 8192 16384 32768 65536 131072 262144 524288 1048576}" ;; + 2) RESOLUTIONS="${RESOLUTIONS:-16 32 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..8a843533 --- /dev/null +++ b/scripts/slurm/submit_forward_time_nd.sh @@ -0,0 +1,238 @@ +#!/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}" + +# Optional kernel override: a host directory of pre-downloaded aarch64 wheels for +# subquadratic-ops-torch-cu13. The image bakes in whatever the [cuda] extra +# resolved at build time (0.2.1 from public PyPI); FFT_BACKEND=subq_ops_fused +# needs >= 0.2.2, which today ships only from the internal NVIDIA GitLab registry. +# Rather than rebuild the .sqsh (no docker on the login node), stage the wheels +# once and pip-install them into the container at job start: +# +# pip download --no-deps --only-binary=:all: \ +# --platform manylinux_2_28_aarch64 --python-version 3.12 --abi cp312 \ +# -d /lustre/.../subq_ops_wheels 'subquadratic-ops-torch-cu13>=0.2.2' \ +# --index-url https://__token__:@gitlab-master.nvidia.com/api/v4/projects/180496/packages/pypi/simple +# SUBQ_OPS_WHEEL_DIR=/lustre/.../subq_ops_wheels sbatch scripts/slurm/submit_forward_time_nd.sh +# +# Installed with --no-index so the compute node needs no network access. +SUBQ_OPS_WHEEL_DIR="${SUBQ_OPS_WHEEL_DIR:-}" +WHEEL_MOUNT=/workspace/subq_wheels +if [[ -n "${SUBQ_OPS_WHEEL_DIR}" ]]; then + if [[ ! -d "${SUBQ_OPS_WHEEL_DIR}" ]]; then + echo "SUBQ_OPS_WHEEL_DIR='${SUBQ_OPS_WHEEL_DIR}' is not a directory"; exit 1 + fi + MOUNTS="${MOUNTS},${SUBQ_OPS_WHEEL_DIR}:${WHEEL_MOUNT}:ro" +fi + +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): every power of two from a 16-wide grid up to ~16M +# tokens. Override with RESOLUTIONS=. +# 1D: L = R 2D: L = R^2 3D: L = R^3 +case "${DATA_DIM}" in + 1) RESOLUTIONS="${RESOLUTIONS:-16 32 64 128 256 512 1024 2048 4096 8192 16384 32768 65536 131072 262144 524288 1048576 2097152 4194304 8388608 16777216}" ;; + 2) RESOLUTIONS="${RESOLUTIONS:-16 32 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 +# subq_ops_fused is the fused 2D cuFFTDx kernel (subquadratic-ops-torch >= 0.2.2). +# It is 2D-only and capped at 64 per axis, so the benchmark resolves it PER POINT: +# 2D R<=64 runs fused, 2D R>=128 and all of 1D fall back to subq_ops, 3D to +# torch_fft. Set FFT_BACKEND=subq_ops for the pre-0.2.2 behaviour everywhere. +FFT_BACKEND="${FFT_BACKEND:-subq_ops_fused}" +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}" +SHORT_CONV="${SHORT_CONV:-subq_ops}" +# Extra flags appended verbatim to the benchmark invocation (submit_nemotron_1d.sh +# uses this for the shared mamba/gdp config knobs). +EXTRA_ARGS="${EXTRA_ARGS:-}" +# The gdp mixer needs flash-linear-attention, which is not in the image. Set +# NEEDS_FLA=1 to pip-install it at job start (compute nodes have network access). +NEEDS_FLA="${NEEDS_FLA:-0}" + +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 | grep -iE "^(torch|triton)==" > /tmp/fla-constraints.txt + pip install --no-cache-dir -c /tmp/fla-constraints.txt flash-linear-attention \ + || { echo "ERROR: flash-linear-attention install failed — gdp cannot run."; exit 1; } + python -c " +import sys +sys.modules.setdefault('tilelang', None) +from fla.ops.gated_delta_product import chunk_gated_delta_product +import importlib.metadata as md +print('fla:', md.version('flash-linear-attention'), '- chunk_gated_delta_product OK') +" || { echo "ERROR: fla installed but the GDP kernel does not import."; exit 1; } +fi + +python -c "import subquadratic_ops_torch; print('subq_ops: available')" \ + || echo "WARNING: subquadratic_ops_torch missing — hyena subq_ops points will error." +python -c " +from importlib.metadata import version +v = version('subquadratic-ops-torch-cu13') +print('subq_ops version:', v) +" || echo "WARNING: could not read subq_ops version." +# subq_ops_fused needs >= 0.2.2; fail loudly instead of silently benchmarking a fallback. +if [ "${FFT_BACKEND}" = "subq_ops_fused" ]; then + python -c " +import sys +from subquadratic_ops_torch.fused_fft_conv2d import fused_fft_conv2d # noqa: F401 +print('subq_ops: fused_fft_conv2d available') +" || { echo "ERROR: FFT_BACKEND=subq_ops_fused but fused_fft_conv2d is missing (needs >= 0.2.2)."; exit 1; } +fi + +PYTHONPATH=. python benchmarks/benchmark_forward_time_nd_resolution.py \ + --data-dim ${DATA_DIM} \ + --fft-backend ${FFT_BACKEND} \ + --short-conv ${SHORT_CONV} \ + ${EXTRA_ARGS} \ + --grid-type ${GRID_TYPE} \ + --dtype ${DTYPE} \ + --batch-size ${BATCH_SIZE} \ + --hidden-dim ${HIDDEN_DIM} \ + --num-heads ${NUM_HEADS} \ + --mamba-headdim ${MAMBA_HEADDIM} \ + --mixers ${MIXERS} \ + --resolutions ${RESOLUTIONS} \ + --num-warmup ${NUM_WARMUP} \ + --num-iters ${NUM_ITERS} \ + --max-seconds-per-point ${MAX_SECONDS} \ + --output ${JSONL} + +# Render the time and memory plots in-job (matplotlib ships in the dev image). +if python -c "import matplotlib" 2>/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/slurm/submit_nemotron_1d.sh b/scripts/slurm/submit_nemotron_1d.sh new file mode 100755 index 00000000..a84fe376 --- /dev/null +++ b/scripts/slurm/submit_nemotron_1d.sh @@ -0,0 +1,95 @@ +#!/bin/bash +# ============================================================================= +# nemotron-1d — HyenaND vs Mamba-2 vs Gated Delta Product at NEMOTRON's config. +# +# A different question from the reach/flash sweeps, and deliberately a different +# setup. Those ask "how far does each operator scale" at a tiny width out to 16.7M +# tokens. This one asks the only question Nemotron's architecture committee will +# care about: **at the config and sequence length Nemotron actually pretrains at, +# how does HyenaND compare to the Mamba-2 it would replace, and to the Gated Delta +# Product that is the competing replacement?** +# +# Why the reach sweep does not answer it: Nemotron pretrains at L=8192, and per +# nemotron_workspace/FLOPS_MATCHING.md the long conv is ~3.4% of a Hyena layer +# there — "the layer is projection-bound". A 6496x advantage at 16.7M says nothing +# about 8192. The comparison here is projections, not asymptotics. +# +# Config — the 1B A315M-88B rung, the decision gate in NEMOTRON_HYENA_PLAN.md: +# hidden 768 | mamba_state_dim 128 | mamba_head_dim 64 | mamba_num_groups 8 +# mamba_num_heads 24 (= d_inner 1536 / 64) | expand 2 | L = 8192 | bf16 +# Sources: NEMOTRON_HYENA_PLAN.md:220,248 and FLOPS_MATCHING.md:125. +# +# GDP reads the SAME mamba_* fields as Mamba-2 and hardcodes num_householder in the +# source (3 in gated_delta_product_original_v4.py, 2 in _nh2), so sizing all three +# from one config is exactly what a drop-in `--spec` swap produces. Its in_proj is +# consequently wider: 10336 vs Mamba-2's 5144 at M=3 (2.01x) — an arithmetic fact +# about GEMM width. Whether that becomes 2x *time* is what this benchmark measures. +# +# Hyena runs FLOP-matched at HYENA_EXPANSION (default e = 2.36), the value +# FLOPS_MATCHING.md derives for parity with a Mamba-2 layer: qkv h -> 3*e*h, long and +# short conv on e*h channels, out_proj e*h -> h. Set HYENA_EXPANSION=1 for evo2's +# default, which is only 0.423x an M layer and therefore NOT a matched comparison. +# +# BOTH Mamba-2 directions are run, as separate series, because the choice flips the +# result: 'mamba' is bidirectional (two scans -- the right vision/ND baseline) and +# 'mamba_causal' is unidirectional (one scan -- what a language model runs, and the +# correct baseline here). Bidirectional does ~2x the work: measured at L=8192 it is +# 1.98x slower at hidden 8 and 2.22x at hidden 768, which is enough to reverse the +# sign of the Hyena comparison. Earlier 1D sweeps compared causal Hyena against +# BIDIRECTIONAL Mamba-2 and so overstated Hyena's advantage by roughly 2x. +# +# The gdp-vs-mamba pairing IS like-for-like: both are sized from the same mamba_* +# config, which is exactly what a drop-in --spec swap produces. +# +# Usage: +# scripts/slurm/submit_nemotron_1d.sh # 1B rung, L sweep +# HIDDEN_DIM=2048 MAMBA_NUM_HEADS=64 scripts/slurm/submit_nemotron_1d.sh # 14B rung +# GDP_HOUSEHOLDER=2 scripts/slurm/submit_nemotron_1d.sh # the _nh2 variant +# SEQ_LENS="8192" scripts/slurm/submit_nemotron_1d.sh # operating point only +# +# Writes benchmarks/results/nemotron_1d.{jsonl,png,pdf} — a NEW output stem, so the +# reach/flash results are untouched. +# ============================================================================= +set -euo pipefail +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +# ── The 1B A315M rung ──────────────────────────────────────────────────────── +HIDDEN_DIM="${HIDDEN_DIM:-768}" +MAMBA_STATE_DIM="${MAMBA_STATE_DIM:-128}" +MAMBA_HEADDIM="${MAMBA_HEADDIM:-64}" +MAMBA_NGROUPS="${MAMBA_NGROUPS:-8}" +MAMBA_EXPAND="${MAMBA_EXPAND:-2}" +GDP_HOUSEHOLDER="${GDP_HOUSEHOLDER:-3}" +# FLOP-matched inner width for Hyena. FLOPS_MATCHING.md derives e=2.36 for parity +# with a Mamba-2 layer; e=1 (evo2 default) is only 0.423x an M layer. +HYENA_EXPANSION="${HYENA_EXPANSION:-2.36}" + +# Sweep around Nemotron's 8192, wide enough to show both the operating point and +# where each operator walls. L = R in 1D. +SEQ_LENS="${SEQ_LENS:-1024 2048 4096 8192 16384 32768 65536 131072 262144}" + +MIXERS="${MIXERS:-hyena mamba_causal mamba gdp}" + +# Attention needs hidden/num_heads to divide by 2 for 1D RoPE; 768/12 = 64. +NUM_HEADS="${NUM_HEADS:-12}" + +DATA_DIM=1 +OUT="${OUT:-nemotron_1d}" +FFT_BACKEND="${FFT_BACKEND:-subq_ops}" # 1D causal fused long conv +SHORT_CONV="${SHORT_CONV:-subq_ops}" # fused causal short conv + +echo "[nemotron-1d] hidden=${HIDDEN_DIM} state_dim=${MAMBA_STATE_DIM} head_dim=${MAMBA_HEADDIM}" +echo "[nemotron-1d] ngroups=${MAMBA_NGROUPS} expand=${MAMBA_EXPAND} gdp_M=${GDP_HOUSEHOLDER}" +echo "[nemotron-1d] mixers='${MIXERS}' L='${SEQ_LENS}' (Nemotron pretrains at 8192)" + +export DATA_DIM MIXERS HIDDEN_DIM NUM_HEADS MAMBA_HEADDIM MAMBA_EXPAND OUT FFT_BACKEND SHORT_CONV +export MAMBA_STATE_DIM MAMBA_NGROUPS GDP_HOUSEHOLDER +export RESOLUTIONS="${SEQ_LENS}" +export EXTRA_ARGS="--mamba-state-dim ${MAMBA_STATE_DIM} --mamba-ngroups ${MAMBA_NGROUPS} --gdp-householder ${GDP_HOUSEHOLDER} --hyena-expansion ${HYENA_EXPANSION}" +export NEEDS_FLA=1 + +exec sbatch \ + --job-name="nvsubq-nemotron-1d" \ + --time="${TIME_LIMIT:-02:00:00}" \ + --export=ALL \ + "${HERE}/submit_forward_time_nd.sh" diff --git a/scripts/visualization/gen_benchmark_digest.py b/scripts/visualization/gen_benchmark_digest.py new file mode 100644 index 00000000..e0d5068c --- /dev/null +++ b/scripts/visualization/gen_benchmark_digest.py @@ -0,0 +1,335 @@ +# 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. + +"""Build a single self-contained results digest from the six forward-time JSONL sweeps. + +Everything numeric is derived from the JSONL rows, not transcribed, so the digest +cannot drift from the data it describes. +""" + +from __future__ import annotations + +import json +from collections import defaultdict +from pathlib import Path + + +# The sweep JSONLs this reads, and the digest it writes, are committed at +# /benchmarks/results/. Derived from this file's location so the +# generator runs from any clone, not just the original Lustre workspace. +RESULTS = Path(__file__).resolve().parents[2] / "benchmarks" / "results" + +SWEEPS = [ + ("forward_time_1d", "Reach 1D", 1, "2559702"), + ("forward_time_2d", "Reach 2D", 2, "2559620"), + ("forward_time_3d", "Reach 3D", 3, "2559621"), + ("forward_time_flash_1d", "Flash-kernel 1D", 1, "2559703"), + ("forward_time_flash_2d", "Flash-kernel 2D", 2, "2559623"), + ("forward_time_flash_3d", "Flash-kernel 3D", 3, "2559624"), +] + +MIXER_ORDER = ["hyena", "attention", "flex", "fa4", "mamba"] + + +def load(stem: str) -> list[dict]: + p = RESULTS / f"{stem}.jsonl" + return [json.loads(line) for line in p.read_text().splitlines() if line.strip()] + + +def cell(row: dict | None) -> str: + """One table cell: the timing, or why the point has none.""" + if row is None: + return "-" + if row["status"] == "ok": + return f"{row['ms']:.3f}" + return {"oom": "oom", "timeout": "timeout", "unavailable": "n/a", "error": "error"}.get( + row["status"], row["status"] + ) + + +def fmt_int(n: int) -> str: + return f"{n:,}" + + +def sweep_section(stem: str, title: str, dim: int, jobid: str) -> list[str]: + rows = load(stem) + if not rows: + return [f"### {title}\n\n(no data)\n"] + + by_mixer_res: dict[tuple[str, int], dict] = {(r["mixer"], r["resolution"]): r for r in rows} + mixers = [m for m in MIXER_ORDER if any(r["mixer"] == m for r in rows)] + resolutions = sorted({r["resolution"] for r in rows}) + meta = rows[0] + + # Which Hyena backends/short-convs actually ran, per resolution. + backends = {r["resolution"]: r.get("backend") for r in rows if r["mixer"] == "hyena"} + short_convs = {r.get("short_conv") for r in rows if r["mixer"] == "hyena"} - {None} + + out = [f"### {title} (`{stem}`)", ""] + out.append(f"- SLURM job: `{jobid}` | device: `{meta['device']}` | dtype: `{meta['dtype']}`") + out.append( + f"- hidden_dim={meta['hidden_dim']}, num_heads={meta['num_heads']}, " + f"batch_size={meta['batch_size']}, data_dim={dim} (L = R^{dim})" + ) + if short_convs: + out.append(f"- Hyena short conv: {', '.join(sorted(short_convs))}") + uniq_backends = sorted({b for b in backends.values() if b}) + if uniq_backends: + detail = [] + for b in uniq_backends: + rs = [r for r, v in backends.items() if v == b] + detail.append(f"`{b}` at R={min(rs)}..{max(rs)}" if len(rs) > 1 else f"`{b}` at R={rs[0]}") + out.append(f"- Hyena FFT backend: {'; '.join(detail)}") + out.append("") + + header = "| R | L | " + " | ".join(mixers) + " |" + sep = "|---|---|" + "---|" * len(mixers) + out += [header, sep] + for R in resolutions: + cells = [cell(by_mixer_res.get((m, R))) for m in mixers] + out.append(f"| {fmt_int(R)} | {fmt_int(R**dim)} | " + " | ".join(cells) + " |") + out.append("") + out.append("All values are ms per forward pass, lower is better.") + out.append("") + + # Largest measured gap vs hyena, computed only where both sides are real numbers. + if "hyena" in mixers: + for other in [m for m in mixers if m != "hyena"]: + best = None + for R in resolutions: + h, o = by_mixer_res.get(("hyena", R)), by_mixer_res.get((other, R)) + if h and o and h["status"] == "ok" and o["status"] == "ok" and h["ms"] > 0: + ratio = o["ms"] / h["ms"] + if best is None or ratio > best[1]: + best = (R, ratio) + if best: + # One decimal below 10x: rounding 1.37x to "1x" reads as "no difference". + r = f"{best[1]:.1f}" if best[1] < 10 else f"{best[1]:.0f}" + out.append(f"- max {other}/hyena: **{r}x** at L={fmt_int(best[0] ** dim)}") + out.append("") + + # Non-ok points, so a reader does not mistake a gap for missing work. + bad = [r for r in rows if r["status"] != "ok"] + if bad: + grouped = defaultdict(list) + for r in bad: + grouped[(r["mixer"], r["status"])].append(r["resolution"]) + out.append("Points without a timing:") + for (m, status), rs in sorted(grouped.items()): + out.append(f"- `{m}` {status} at R={sorted(rs)}") + out.append("") + return out + + +def iters_note() -> list[str]: + """Show how the adaptive iteration count varies — it matters for reading the tail.""" + rows = [r for r in load("forward_time_1d") if r["mixer"] == "attention" and r["status"] == "ok"] + out = ["| L | iterations |", "|---|---|"] + for r in sorted(rows, key=lambda r: r["seq_len"]): + out.append(f"| {fmt_int(r['seq_len'])} | {r.get('iters', '?')} |") + return out + + +def main() -> None: + doc: list[str] = [] + doc += [ + "# nvSubquadratic forward-time benchmarks — results digest", + "", + "Single-layer forward-pass timings for HyenaND vs attention kernels vs Mamba2,", + "swept from a 16-wide grid up to ~16.7M tokens in 1D, 2D and 3D.", + "", + "This file is generated from the JSONL sweep outputs in this directory; every number", + "below is read from that data rather than transcribed. Raw per-point records live in", + "`forward_time_{1,2,3}d.jsonl` and `forward_time_flash_{1,2,3}d.jsonl`, with matching", + "`.png`/`.pdf` plots and `*-.out` run logs.", + "", + "---", + "", + "## What was run", + "", + "Two sweep configurations, each across 1D/2D/3D — six jobs total. They answer", + "different questions and are **not** comparable to each other:", + "", + "| Config | hidden_dim | head_dim | Mixers | Reaches | Purpose |", + "|---|---|---|---|---|---|", + "| **Reach** | 8 | 4 | hyena, attention, mamba | 16.7M tokens | How far each operator scales. Small width keeps the qkv tensor under torch's 2^31 index limit at 16.7M. |", + "| **Flash-kernel** | 512 | 128 | hyena, attention, flex, fa4, mamba | ~1M tokens | SDPA vs FlexAttention vs FlashAttention-4 vs HyenaND at the width flash kernels are optimised for. Walls at ~1.4M on the same index limit. |", + "", + "`attention`, `flex` and `fa4` are three interchangeable kernels on one shared q/k/v +", + "RoPE path. `flex` and `fa4` require head_dim >= 16, so they cannot run in the Reach", + "config and appear only in the Flash-kernel sweeps.", + "", + "---", + "", + "## Environment", + "", + "| Component | Version |", + "|---|---|", + "| GPU | NVIDIA GB200 (aarch64), driver 580.173.02 |", + "| torch | 2.12.1+cu130 (CUDA 13.0) |", + "| subquadratic-ops-torch-cu13 | 0.2.2 |", + "| mamba-ssm / causal-conv1d | 2.3.2.post1 / 1.6.2.post1 |", + "| flash-attn-4 | 4.0.0b23 (cutlass-dsl 4.6.0.dev0) |", + "| apex | 0.1 |", + "| Image | `/lustre/fsw/healthcareeng_bionemo/farhadr/enroot/nvsubquadratic-arm64.sqsh` |", + "| Branch | `farhadr/2d_bench` (merge of PR #137 + PR #138) |", + "", + "---", + "", + "## Results", + "", + ] + + for stem, title, dim, jobid in SWEEPS: + doc += sweep_section(stem, title, dim, jobid) + doc.append("---") + doc.append("") + + doc += [ + "## How to read these numbers", + "", + "### 1. Below ~65K tokens the sweep does not measure attention", + "", + "In the Reach config (head_dim 4) attention sits flat at ~1.1 ms from R=32 to R=32768.", + "That floor is fixed per-iteration dispatch overhead, not attention: the quadratic term", + "is real but far below the floor until L=65536, where attention first departs (5.02 ms)", + "and then scales cleanly (17.2 -> 67.8 -> 262 -> 1047 ms, ~4x per doubling).", + "", + "**Do not draw conclusions from points below L=65536.** For a figure, truncating the", + "x-axis there is the honest presentation.", + "", + "### 2. There is a reproducible dip at R=2048 in the 1D Reach plot", + "", + "Investigated directly. It is **not** hardware and **not** a one-off:", + "", + "- Clean ECC (0 uncorrected, all 4 GPUs), no throttling of any kind, 27 C.", + "- Reproduces on a different node (`ptyche0178`) than the production run (`ptyche0077`),", + " ruling out a node-specific fault.", + "- 6 repetitions at R=2048: min 0.585, median 0.649, max 1.189 ms. Five of six land at", + " ~0.65 ms, so the production value of 0.643 is the *typical* reading, not an outlier.", + "", + "Timings in this region are **bimodal**, landing near ~0.6 ms or ~1.2 ms, and which", + "resolutions land where changes with node and process context (in production only R=16", + "and R=2048 were fast; on the rerun node R=1024 and R=2048 were fast while R=4096 was", + "slow). The likely mechanism is SDPA backend selection: at head_dim 4 the flash path is", + "ineligible, so PyTorch chooses between math and mem-efficient kernels by heuristic,", + "giving two stable states. Re-running relocates the dip rather than removing it.", + "", + "### 3. Mamba's `error` points are kernel limits, NOT out-of-memory", + "", + "No point in any sweep ran out of memory. The `error` entries here are", + "`cudaErrorInvalidValue` raised by Dao-AILab's `causal_conv1d`: its channels-last path", + "tiles the sequence at 64 elements per block and hits the CUDA grid-dimension cap of", + "65,535, giving a hard ceiling of **4,194,240 tokens**. That is a launch-configuration", + "limit, so it occurs on any GPU regardless of memory (peak use here was under 5 GB).", + "", + "At wider configurations a *different* and earlier limit binds first: 32-bit index", + "overflow in the Triton SSD scan (`ssd_chunk_state.py`), which fails at ~2^31 elements", + "and therefore moves inversely with model width (~600K tokens at hidden 768). These", + "sweeps use hidden_dim 8, which is why they reach 4M before failing.", + "", + "Both limits are implementation properties of the shipping kernels, not architectural", + "bounds on Mamba-2, and both are documented with reproductions in", + "`docs/mamba2_limits.md`. The Mamba baseline intentionally keeps its own", + "`causal_conv1d`; it was not switched to the subq_ops kernel.", + "", + "### 4. The fused 2D kernel covers only part of the 2D sweep", + "", + "`fft_backend=subq_ops_fused` (subq_ops >= 0.2.2) is **2D-only and capped at 64 per", + "axis** — its largest FFT tile is 128 and it requires max(X,Y) <= fft_size/2. Across a", + "16..16M sweep it therefore covers 2D at R=16/32/64 only; R>=128 falls back to", + "`subq_ops`, and 3D to `torch_fft` (no 3D CUDA kernel exists). The backend actually used", + "is recorded per row in the JSONL as `backend`, so no point is misattributed.", + "", + "### 5. 1D HyenaND is not comparable to runs before 2026-08-10", + "", + "HyenaND's short conv now follows the operator's causality. Previously the causal 1D", + "config paired a causal long conv with a symmetric-padded short conv, so the operator", + "could see one token of future context. 1D now uses the fused", + "`subquadratic_ops_torch.causal_conv1d` (left-only padding), recorded per row as", + "`short_conv`. This is a behaviour change, not only a speedup. 2D/3D are unaffected:", + "they are non-causal and still use symmetric `torch.nn.ConvNd`, because the fused kernel", + "is causal and 1D-only.", + "", + "### 6. The Flash-kernel 3D sweep is too short to show a crossover", + "", + "At hidden_dim 512 in 3D the 2^31 index limit bites at R=64 (262K tokens), leaving only", + "three points. HyenaND is still *behind* attention at the first two and only pulls ahead", + "at the last (70.9 vs 96.8 ms, 1.4x). Treat that sweep as a kernel-availability check", + "rather than evidence about scaling; the 3D scaling story is in the Reach 3D sweep,", + "which reaches 16.7M.", + "", + "### 7. Points in the tail are averaged over very few iterations", + "", + "Each datapoint is a single measurement: CUDA events around one back-to-back loop of N", + "forwards, divided by N. It is **not** an average of repeated runs. N adapts to the cost", + "(target ~5 s per point, capped by `--num-iters` and `--max-seconds-per-point`), so the", + "largest points have very few iterations — the 16.7M attention point is N=1. Those are", + "still reliable because a 268-second measurement swamps jitter, but they are single-shot.", + "", + "1D attention iteration counts:", + "", + ] + doc += iters_note() + doc += [ + "", + "---", + "", + "## JSONL schema", + "", + "One record per (mixer, resolution):", + "", + "```json", + json.dumps(load("forward_time_1d")[0], indent=2), + "```", + "", + "| Field | Meaning |", + "|---|---|", + "| `status` | `ok`, `oom`, `timeout`, `error`, or `unavailable` (dependency missing). Only `ok` rows have `ms`. |", + "| `ms` | Milliseconds per forward pass. |", + "| `iters` | Iterations the timed loop averaged over (adaptive; see note 6). |", + "| `backend` | Hyena FFT-conv backend actually used at this point (null for non-hyena). |", + "| `short_conv` | Hyena short-conv implementation actually used (null for non-hyena). |", + "| `mem_gb` | Peak CUDA memory allocated during the point. |", + "", + "## Reproducing", + "", + "```bash", + "# Reach sweeps (16 -> 16.7M):", + "for D in 1 2 3; do DATA_DIM=$D sbatch --time=04:30:00 --export=ALL \\", + " scripts/slurm/submit_forward_time_nd.sh; done", + "", + "# Flash-kernel sweeps:", + "for D in 1 2 3; do DATA_DIM=$D scripts/slurm/submit_forward_time_flash_kernels.sh; done", + "```", + "", + "Both wrap `benchmarks/benchmark_forward_time_nd_resolution.py`. Relevant flags:", + "`--fft-backend {subq_ops,subq_ops_fused,torch_fft}`, `--short-conv {subq_ops,torch}`,", + "`--mixers`, `--resolutions`, `--hidden-dim`, `--num-iters`, `--max-seconds-per-point`.", + "", + ] + + out_path = RESULTS / "BENCHMARK_RESULTS.md" + out_path.write_text("\n".join(doc) + "\n") + print(f"wrote {out_path} ({len(doc)} lines, {out_path.stat().st_size / 1024:.1f} KB)") + # The tables above are emitted unaligned ("| a | b |") and square brackets + # are left unescaped, both of which mdformat rewrites. The committed file is + # therefore NOT byte-identical to this output, and regenerating it will fail + # the lint job until the hook is re-run: + print(f"NOTE: run `pre-commit run --files {out_path}` before committing — mdformat will reformat it.") + + +if __name__ == "__main__": + main() diff --git a/scripts/visualization/visualize_forward_time_nd.py b/scripts/visualization/visualize_forward_time_nd.py new file mode 100644 index 00000000..ffa77c8b --- /dev/null +++ b/scripts/visualization/visualize_forward_time_nd.py @@ -0,0 +1,419 @@ +# 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} + if not ok_seq: + raise SystemExit(f"No successful ('ok') {mkey} values in the JSONL — nothing to plot.") + 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 / error / 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()