Mixer Benchmarks - #137
Open
farhadrgh wants to merge 33 commits into
Open
Conversation
Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com>
Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com>
Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com>
Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com>
Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com>
Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com>
Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com>
Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com>
Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com>
Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com>
Signed-off-by: Farhad Ramezanghorbani <farhadr@nvidia.com>
JSONL data and time/memory plots for the GB200 forward-time sweeps: - results_8_hidden_16M_tokens/: hidden 8 reach sweep (HyenaND / Attention / Mamba2 out to 16M tokens; attention is the non-flash head_dim-4 path). - results_512_hidden/: hidden 512, head_dim 128 flash-kernel comparison (HyenaND vs SDPA / FlexAttention / FlashAttention-4 / Mamba2). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Consolidate the two benchmark sets under benchmarks/results/ and document them in one README (protocol, configs, headline numbers, reproduce steps): - reach_hidden8/ (hidden 8, head_dim 4 — scaling reach to 16M tokens) - flash_hidden512/ (hidden 512, head_dim 128 — SDPA/Flex/FA4/Mamba2 comparison) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Adds `subquadratic_ops_torch.fused_fft_conv2d` as a first-class 2D FFT-conv path. Unlike every other FFT backend it runs the whole rfft2/multiply/irfft2 pipeline in a single cuFFTDx launch and natively in fp32/fp16/bf16 rather than upcasting to fp32. Measured on H200 (B=8, H=768, fwd+bwd, bf16): 3.6-3.9x over `torch_fft` and 1.2-2.4x over the existing `subq_ops` path. Two entry points: * `fft_backend="subq_ops_fused"` on CKConvND — explicit and predictable. Restricted to data_dim=2, non-causal, zero padding, and spatial extents of at most 64 per axis (the kernel's largest FFT tile is 128 and it requires max(X, Y) <= fft_size // 2). The spatial cap is enforced on the first forward pass, not at construction, since input size is unknown there. * `nvsubquadratic.ops.fftconv_lowering` — an inductor pre-grad pass that rewrites fftconv.py's chain onto the fused kernel, so a model already on `fft_backend="torch_fft"` picks it up without a config change. This matters because inductor cannot codegen complex operators and otherwise falls back to eager cuFFT for the entire chain. Enable per-callable with `torch.compile(model, options=fused_fftconv2d_options())`, or for a scope with the `fused_fftconv2d_lowering()` context manager. Pre-grad rather than post-grad so autograd is derived from the custom op's registered backward instead of requiring a consistent forward+backward rewrite. It fires only on an exact match of the reference recipe (padding rule, crop offset, shape limits, CUDA device, and a compute capability that supports the required tile — the 128 tile needs SM90+, as SM80/SM86 lack the shared memory). `lowering_stats()` reports rewrite and per-reason skip counts, since a silent pass is otherwise indistinguishable from one that never ran. Crop-offset reconciliation: the upstream kernel crops the 'same' window at fft_size // 2 whereas fftconv.py crops at K // 2. The wrapper pre-pads the filter's top/left by the difference, making results interchangeable with the other backends (~3e-7 normwise in fp32, ~3e-3 in bf16). Without it the output is shifted by fft_size // 2 - K // 2 pixels — which reads as a ~1.41 (sqrt 2, fully decorrelated) error, not as reduced accuracy. A regression test pins this so the pre-pad cannot be optimised away. Also registers the fused operators eagerly when the pass is constructed: inductor's on-disk FX cache lets a compiled artifact call torch.ops.subquadratic_ops_torch.* on a cache hit without any Python call having triggered the wrappers' lazy import, which failed with an opaque op-namespace AttributeError. Fixes a pre-existing test-gate bug: tests/conftest.py resolved the kernel version from the `subquadratic-ops-torch-cu12` distribution only. On the `-cu13` install that pyproject.toml pins it returned (0, 0, 0), silently turning `requires_subq_ops_v2` into a blanket xfail and hiding every `subq_ops` test result. It now checks both distributions. Includes the in-flight CUDA 13.2 build changes already present in the working tree (Dockerfile base image, cu132 torch/DALI pins, cu13 kernel distribution). Tests: 138 new (op-level equivalence across dtypes/shapes/layouts/FiLM, forward and backward, CKConvND integration, and lowering fire/decline behaviour). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
- Dockerfile: base nvcr.io/nvidia/cuda:13.0.3-devel (matched to torch cu130 so apex/mamba's CUDA-version check passes), torch/torchvision cu130, nvidia-dali-cuda130. The [cuda] extra resolves subquadratic-ops-torch-cu13 (>=0.2.2) from the internal GitLab registry via a build secret; FlashAttention-4 uses the [cu13] extra. - build_sqsh.sh: pass the gitlab_token secret (+ preflight), and refresh QEMU binfmt for arm64-on-x86 cross-builds (fixes the nvcc SIGSEGV under stale QEMU). - pyproject: [cuda] → subquadratic-ops-torch-cu13, [dali] → nvidia-dali-cuda130. - Docs / examples / conftest / CHANGELOG updated to CUDA 13.0 / cu130 / cu13. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Brings in fft_backend="subq_ops_fused" and the torch.compile lowering pass so the 2D forward-time benchmarks can exercise the fused cuFFTDx kernel. Conflict resolutions: * pyproject.toml — take #138's subquadratic-ops-torch-cu13>=0.2.2 pin (needed for fused_fft_conv2d), noting that 0.2.2 currently ships only from the internal NVIDIA GitLab registry (public PyPI tops out at 0.2.1). * Dockerfile — keep this branch's CUDA 13.0 base + torch 2.10.0/cu130 pins rather than #138's 13.2/2.12.1: apex and mamba only build when the base nvcc CUDA matches torch's exactly, and the benchmark image builds both. Adopted #138's ARG parameterisation with cu130 defaults so a cu132 image is a build-arg away, and added SUBQ_OPS_INDEX_URL for the 0.2.2 registry. * tests/conftest.py — take #138's version (queries both the cu12 and cu13 distributions instead of only cu13). * CHANGELOG.md — keep both entries. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
benchmark_forward_time_nd_resolution.py
* --fft-backend gains "subq_ops_fused".
* resolve_hyena_backend() picks the effective backend PER POINT instead of
once per run. The fused kernel is 2D-only and capped at 64 per axis (its
largest FFT tile is 128 and it needs max(X,Y) <= fft_size//2), so a single
--fft-backend choice cannot cover a 16..16M sweep: 2D R<=64 runs fused, 2D
R>=128 and all of 1D fall back to subq_ops, 3D to torch_fft. The per-point
backend is echoed before the sweep and recorded in each JSONL row, so the
plots do not imply the fused kernel ran where it could not.
submit_forward_time_nd.sh
* Sweeps every power of two from a 16-wide grid to ~16M tokens (1D 16..16M,
2D 16..4096, 3D 16..256).
* FFT_BACKEND now defaults to subq_ops_fused.
* SUBQ_OPS_WHEEL_DIR mounts a host directory of pre-staged wheels and installs
them offline (--no-index) at job start, so the newer kernel can be swapped in
without rebuilding the .sqsh. Verifies fused_fft_conv2d imports and aborts if
it does not, rather than silently timing the fallback path.
submit_forward_time_flash_kernels.sh
* Extend the same sweeps down to R=16.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The merge pinned subquadratic-ops-torch-cu13>=0.2.2 for fused_fft_conv2d, but public PyPI only carries 0.2.1, so the Dockerfile's `.[all]` layer now fails on a clean build. Resolve it from an internal GitLab index instead. The token is passed as a BuildKit secret rather than a --build-arg: a build-arg is recorded in the image's layer history, so the token would ship with every .sqsh built from it. build_sqsh.sh reads GITLAB_TOKEN from the environment or ~/.gitlab_token, builds the index URL, and fails early with the token-creation steps rather than letting the build run for hours before dying on pip install. Drop the secret once 0.2.2 is published to public PyPI. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
build_sqsh.sh needs `docker buildx`, plus QEMU binfmt emulation to produce an arm64 image from an x86 host. Cluster login nodes here have neither — no docker daemon or any other builder, no qemu-aarch64 registered, and no root to install either — so there was no way to produce a GB200 image from the login node. `enroot import` is not a way around it: /tmp is a 2 GB tmpfs (too small to extract a CUDA devel image) and redirecting it to lustre fails because lustre cannot hold the capability xattrs enroot sets while extracting layers. Instead, let pyxis pull the base image onto a GB200 node, take root inside it with --container-remap-root, replay the Dockerfile's steps natively, and write the result out with --container-save. Building on the target architecture also removes QEMU, so the MAX_JOBS=1 throttle build_sqsh.sh needs to survive emulated apex/mamba compiles does not apply: this defaults to MAX_JOBS=32. The token reaches the container as a read-only mounted file rather than an env var — srun's environment is visible via `scontrol show job`, and an exported var would be captured into the saved image. This script REPLAYS the Dockerfile rather than parsing it, so the version pins are duplicated in both and will drift if only one is edited. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
PR #138 raised pyproject's floor to torch>=2.12.0,<2.13.0, but the Dockerfile kept TORCH_VERSION=2.10.0. Those two contradict, and the failure was silent: apex, mamba-ssm and causal-conv1d compiled against 2.10.0 in the middle of the build, then the final `.[all]` step resolved nvsubquadratic's own requirement and upgraded torch to 2.12.1 underneath them — leaving compiled extensions built against headers that no longer matched the installed torch. Nothing caught this, because the build's verification printed package METADATA versions, which report happily regardless of what the extension was built against. Add a guard that compares torch.__version__ to the pin and fails the build, and note the coupling on both pins so they are not edited apart. Also fix the FA4 probe in both files: nvidia_cutlass_dsl exposes no __version__, so reading the attribute raised AttributeError and made a perfectly working FA4 install report as failed on every build. Read the version from metadata instead. Verified on the rebuilt image: torch 2.12.1+cu130, apex / causal_conv1d / flash_attn.cute / Mamba2 / fused_fft_conv2d all import, and causal_conv1d_fn executes on-GPU. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
HyenaND's short conv was always a torch ConvNd with symmetric padding, in every
config. SubqOpsCausalConv1d — the wrapper around the fused
subquadratic_ops_torch.causal_conv1d kernel — existed and was unit-tested, but
nothing on any Hyena path used it (the module docstring cites two example configs
that in fact select CausalConv1D / torch.nn.Conv1d).
Make the short conv follow the operator's causality rather than always padding
symmetrically:
* is_causal (the 1D sweeps) -> left-only padding: SubqOpsCausalConv1d with
--short-conv=subq_ops (the default), else CausalConv1D. The kernel is
depthwise-only, which this config already satisfies (groups == in_channels
== out_channels == 3*hidden_dim), stride/dilation 1.
* non-causal (2D/3D) -> symmetric torch.nn.ConvNd, unchanged. A causal 1D
kernel does not apply there, so --short-conv=subq_ops falls back with a note.
This also closes a causality hole in the 1D config: it paired a causal long conv
with a symmetric short conv, so the operator could see one token of future
context. 1D Hyena timings will therefore shift slightly against earlier runs.
The effective short conv is printed before the sweep and recorded per JSONL row,
so a plot cannot silently mix short-conv implementations across points.
Mamba keeps its own causal_conv1d (Dao-AILab's) — untouched.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The fused kernel picks its specialisation from the INPUT dtype and then requires
every tensor to match exactly. Under torch.autocast the activations arrive as
bf16/fp16 while nn.Module parameters stay fp32, so passing self.weight through
unchanged aborted with:
ValueError: in_w expected dtype (code=4, bits=16) but got (code=2, bits=32)
That made the module unusable on the standard mixed-precision path — it only
worked when the caller had already narrowed the parameters by hand, which is why
the existing tests (all same-dtype) did not catch it.
Cast weight/bias to the input dtype in forward, mirroring what autocast does for
the built-in conv ops: fp32 master parameters are untouched, only the values
handed to the kernel are narrowed.
Verified on GB200 against CausalConv1D: fp32 max|fused-ref| = 2.4e-07, bf16
1.6e-02, fp16 2.0e-03 (rounding), and perturbing position t changes outputs at
t..L-1 while leaving 0..t-1 bit-identical.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Every existing test feeds fp32 activations to fp32 parameters, so all tensors agree by construction and the kernel's exact-dtype check passes. That is why the module could not run under torch.autocast at all without any test noticing. Add TestAutocast: that the module runs under bf16/fp16 autocast, that the fp32 master weight is not mutated (only the kernel's inputs are narrowed), and that the result still matches CausalConv1D at a tolerance appropriate to the narrowed dtype rather than the fp32 ATOL. 14 passed on GB200. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Environment setup
Create the conda environment (required to run tests):
Test plan
pre-commit run --all-filespasses (pre-commit installif not yet set up).pytest tests/).Documentation checklist
For every new or modified public symbol in
nvsubquadratic/orexperiments/:Args:andReturns:blocks with tensor shapes where applicable.r"""..."""(required by ruff D301).docs-tracker.mdwith status[x].