Kernels:

shi-labs
/

natten

Trusted publisher

Uploaded using `kernel-builder`.

by kernels-bot - opened 15 days ago

base: refs/heads/v1

←

from: refs/pr/8

Discussion Files changed

+254315

-0

This view is limited to 50 files because it contains too many changes. See the raw diff here.

Files changed (50) hide show

README.md +112 -0
build/torch210-cu126-aarch64-linux/__init__.py +180 -0
build/torch210-cu126-aarch64-linux/_environment.py +59 -0
build/torch210-cu126-aarch64-linux/_libnatten/__init__.py +109 -0
build/torch210-cu126-aarch64-linux/_libnatten/torch_wrappers.py +1006 -0
build/torch210-cu126-aarch64-linux/_natten_cuda_5fe9dbb.abi3.so +3 -0
build/torch210-cu126-aarch64-linux/_ops.py +9 -0
build/torch210-cu126-aarch64-linux/_private_for_testing.py +33 -0
build/torch210-cu126-aarch64-linux/_types.py +85 -0
build/torch210-cu126-aarch64-linux/attn_merge.py +292 -0
build/torch210-cu126-aarch64-linux/backends/__init__.py +254 -0
build/torch210-cu126-aarch64-linux/backends/blackwell_fmha.py +254 -0
build/torch210-cu126-aarch64-linux/backends/blackwell_fna.py +500 -0
build/torch210-cu126-aarch64-linux/backends/configs/__init__.py +584 -0
build/torch210-cu126-aarch64-linux/backends/configs/checks.py +750 -0
build/torch210-cu126-aarch64-linux/backends/configs/cutlass/__init__.py +422 -0
build/torch210-cu126-aarch64-linux/backends/configs/cutlass/backward_knobs.py +228 -0
build/torch210-cu126-aarch64-linux/backends/configs/cutlass/fna_backward_128x128.py +304 -0
build/torch210-cu126-aarch64-linux/backends/configs/cutlass/fna_backward_128x64.py +223 -0
build/torch210-cu126-aarch64-linux/backends/configs/cutlass/fna_backward_64x64.py +168 -0
build/torch210-cu126-aarch64-linux/backends/configs/cutlass/fna_forward_32x128.py +90 -0
build/torch210-cu126-aarch64-linux/backends/configs/cutlass/fna_forward_64x128.py +82 -0
build/torch210-cu126-aarch64-linux/backends/configs/cutlass/fna_forward_64x64.py +63 -0
build/torch210-cu126-aarch64-linux/backends/configs/cutlass_blackwell/__init__.py +391 -0
build/torch210-cu126-aarch64-linux/backends/configs/cutlass_hopper/__init__.py +522 -0
build/torch210-cu126-aarch64-linux/backends/configs/flex/__init__.py +210 -0
build/torch210-cu126-aarch64-linux/backends/flex.py +799 -0
build/torch210-cu126-aarch64-linux/backends/fmha.py +283 -0
build/torch210-cu126-aarch64-linux/backends/fna.py +417 -0
build/torch210-cu126-aarch64-linux/backends/hopper_fmha.py +261 -0
build/torch210-cu126-aarch64-linux/backends/hopper_fna.py +512 -0
build/torch210-cu126-aarch64-linux/backends/reference.py +343 -0
build/torch210-cu126-aarch64-linux/context.py +231 -0
build/torch210-cu126-aarch64-linux/functional.py +1151 -0
build/torch210-cu126-aarch64-linux/metadata.json +84 -0
build/torch210-cu126-aarch64-linux/metadata.json.sigstore +1 -0
build/torch210-cu126-aarch64-linux/modules.py +449 -0
build/torch210-cu126-aarch64-linux/symbols.json +1257 -0
build/torch210-cu126-aarch64-linux/token_permute/__init__.py +32 -0
build/torch210-cu126-aarch64-linux/token_permute/cutlass_impl.py +286 -0
build/torch210-cu126-aarch64-linux/token_permute/frontend.py +137 -0
build/torch210-cu126-aarch64-linux/token_permute/torch_impl.py +368 -0
build/torch210-cu126-aarch64-linux/utils/__init__.py +22 -0
build/torch210-cu126-aarch64-linux/utils/checks.py +726 -0
build/torch210-cu126-aarch64-linux/utils/device.py +50 -0
build/torch210-cu126-aarch64-linux/utils/dtype.py +36 -0
build/torch210-cu126-aarch64-linux/utils/environment.py +79 -0
build/torch210-cu126-aarch64-linux/utils/log.py +134 -0
build/torch210-cu126-aarch64-linux/utils/tensor.py +113 -0
build/torch210-cu126-aarch64-linux/utils/testing.py +149 -0

README.md ADDED Viewed

	@@ -0,0 +1,112 @@

+---
+license: mit
+tags:
+- kernel
+---
+![Status](https://hubwebhook.dholtz.com/shield?repo=kernels-community/natten)
+# NATTEN — Neighborhood Attention kernels
+CUDA kernels for [NATTEN](https://natten.org) (Neighborhood Attention
+Extension, [SHI-Labs/NATTEN](https://github.com/SHI-Labs/NATTEN)), packaged
+for the [kernels](https://github.com/huggingface/kernels) library. Vendored
+from NATTEN v0.21.7 (MIT license).
+Neighborhood attention restricts each token's attention to a sliding local
+window over 1-D, 2-D, or 3-D token layouts, with optional stride, dilation,
+and per-dimension causal masking. This kernel ships NATTEN's fused
+neighborhood attention (FNA) and FMHA kernels along with the thin functional
+frontend (`na1d`, `na2d`, `na3d`, `attention`, `merge_attentions`) and its
+backend/config selection logic.
+## Usage
+```python
+import torch
+from kernels import get_kernel
+natten = get_kernel("kernels-community/natten", version=1)
+# 2-D neighborhood attention over a 32x32 token layout:
+# [batch, *token_layout, heads, head_dim]
+q = torch.randn(1, 32, 32, 8, 64, device="cuda", dtype=torch.bfloat16)
+k = torch.randn_like(q)
+v = torch.randn_like(q)
+out = natten.functional.na2d(q, k, v, kernel_size=(7, 7))
+# 3-D, with stride/dilation/causal masking per dim:
+q3 = torch.randn(1, 8, 16, 16, 8, 64, device="cuda", dtype=torch.bfloat16)
+k3, v3 = torch.randn_like(q3), torch.randn_like(q3)
+out3 = natten.functional.na3d(
+    q3, k3, v3, kernel_size=(3, 5, 5), dilation=(1, 2, 2), is_causal=(True, False, False)
+)
+```
+All functions are differentiable and torch.compile-compatible (the underlying
+ops are registered via `TORCH_LIBRARY` with fake/meta implementations).
+## Backend / architecture support
+| Backend | Compute capabilities | Notes |
+|---|---|---|
+| Reference CUDA kernels | 7.0 – 12.0 | correctness fallback, all dims |
+| CUTLASS 2.X FNA / FMHA (`cutlass-fna`) | 7.0 – 12.0 (SM80 tensor-core paths) | default on pre-Hopper |
+| Hopper FNA / FMHA (`hopper-fna`) | 9.0a | H100/H200 |
+| Blackwell FNA / FMHA (`blackwell-fna`) | 10.0a | B200/GB200; requires CUDA ≥ 12.8, so cu126 builds fall back to `cutlass-fna` / reference on these GPUs |
+| Flex Attention backend (`flex-fna`) | any | pure PyTorch, via `torch.nn.attention.flex_attention` |
+Known gaps vs. upstream wheels:
+- No SM103 (B300 / Blackwell Ultra) binaries — kernel-builder does not
+  currently target `10.3a`.
+- SM120 (RTX 50 series) gets the CUTLASS 2.X and reference paths only, same
+  as upstream (NATTEN has no SM120-specific fused kernels).
+- No ROCm or CPU support (upstream dropped these in v0.20).
+The backend is chosen automatically per GPU; pass `backend="..."` to
+`na1d`/`na2d`/`na3d` to override.
+## Developing
+The ~144 kernel instantiation TUs under `csrc/autogen/` are committed
+codegen output. To regenerate (e.g. when bumping the vendored NATTEN
+version), run:
+```bash
+scripts/regen.sh                 # re-runs autogen + rewrites build.toml src lists
+```
+Build and test with kernel-builder:
+```bash
+nix run .#build-and-copy -L
+nix run .#ci-test -L
+```
+## Credits
+All kernels and the Python frontend are by Ali Hassani and the NATTEN
+contributors ([SHI-Labs/NATTEN](https://github.com/SHI-Labs/NATTEN), MIT).
+This repository only repackages them for the Kernel Hub: pybind11 bindings
+are replaced with `TORCH_LIBRARY` registrations (Python limited API), and
+build-time codegen is committed statically.
+If you use NATTEN, please cite:
+```bibtex
+@inproceedings{hassani2023neighborhood,
+  title   = {Neighborhood Attention Transformer},
+  author  = {Ali Hassani and Steven Walton and Jiachen Li and Shen Li and Humphrey Shi},
+  year    = 2023,
+  booktitle = {IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
+}
+@misc{hassani2024faster,
+  title   = {Faster Neighborhood Attention: Reducing the O(n^2) Cost of Self Attention at the Threadblock Level},
+  author  = {Ali Hassani and Wen-Mei Hwu and Humphrey Shi},
+  year    = 2024,
+  eprint  = {2403.04690},
+  archivePrefix = {arXiv},
+}
+```

build/torch210-cu126-aarch64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,180 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+from ._environment import HAS_LIBNATTEN
+from .backends import (
+    get_bwd_configs_for_cutlass_blackwell_fmha,
+    get_bwd_configs_for_cutlass_blackwell_fna,
+    get_bwd_configs_for_cutlass_fmha,
+    get_bwd_configs_for_cutlass_fna,
+    get_bwd_configs_for_cutlass_hopper_fmha,
+    get_bwd_configs_for_cutlass_hopper_fna,
+    get_configs_for_cutlass_blackwell_fmha,
+    get_configs_for_cutlass_blackwell_fna,
+    get_configs_for_cutlass_fmha,
+    get_configs_for_cutlass_fna,
+    get_configs_for_cutlass_hopper_fmha,
+    get_configs_for_cutlass_hopper_fna,
+    get_configs_for_flex_fmha,
+    get_configs_for_flex_fna,
+)
+from .context import (
+    allow_flex_compile,
+    allow_flex_compile_backprop,
+    are_deterministic_algorithms_enabled,
+    disable_flex_compile,
+    disable_flex_compile_backprop,
+    get_memory_usage_preference,
+    is_flex_compile_allowed,
+    is_flex_compile_backprop_allowed,
+    is_kv_parallelism_in_fused_na_enabled,
+    is_memory_usage_default,
+    is_memory_usage_strict,
+    is_memory_usage_unrestricted,
+    set_memory_usage_preference,
+    use_deterministic_algorithms,
+    use_kv_parallelism_in_fused_na,
+)
+from .functional import attention, merge_attentions, na1d, na2d, na3d
+from .modules import (
+    NeighborhoodAttention1D,
+    NeighborhoodAttention2D,
+    NeighborhoodAttention3D,
+)
+from . import _private_for_testing  # noqa: F401
+from .version import __version__
+# kernel-builder port: the package contents are installed flat into the build
+# variant directory, so a module literally named `types` would shadow the
+# standard library `types` module whenever that directory is on PYTHONPATH
+# (e.g. kernel-builder test shells and CI runners), breaking interpreter
+# startup. The module therefore lives in `_types`; alias it here so
+# `natten.types` keeps working like upstream.
+import sys as _sys
+from . import _types as types
+_sys.modules[__name__ + ".types"] = types
+# kernel-builder's compat shim (`natten/__init__.py` inside the build variant
+# directory) executes this package under a path-derived module name and copies
+# our globals into a `natten` module whose __path__ contains no submodules.
+# Attribute access (`natten.functional`) works there, but real submodule
+# imports (`from natten.functional import na2d`, `import natten.utils.testing`)
+# would either fail with ModuleNotFoundError or — when resolvable through a
+# parent package's __path__ — re-execute the module under a second name,
+# duplicating module state. Bridge this with a meta-path finder that resolves
+# any `natten.*` import to our already-loaded module objects. Only installed
+# when `natten` in sys.modules is *our* compat shim, so a real `natten`
+# distribution in the same environment is never hijacked.
+if __name__ != "natten":
+    from pathlib import Path as _Path
+    _compat = _sys.modules.get("natten")
+    _is_our_compat = (
+        _compat is not None
+        and getattr(_compat, "__file__", None) is not None
+        and _Path(_compat.__file__).resolve()
+        == _Path(__file__).resolve().parent / "natten" / "__init__.py"
+    )
+    if _is_our_compat:
+        import importlib as _importlib
+        from importlib.abc import Loader as _Loader
+        from importlib.abc import MetaPathFinder as _MetaPathFinder
+        from importlib.util import spec_from_loader as _spec_from_loader
+        _real_root = __name__
+        class _NattenAliasLoader(_Loader):
+            def __init__(self, module):
+                self._module = module
+                self._spec = getattr(module, "__spec__", None)
+                self._loader = getattr(module, "__loader__", None)
+            def create_module(self, spec):
+                return self._module
+            def exec_module(self, module):
+                # The import machinery stamped the alias spec onto the real
+                # module in module_from_spec; restore its original identity.
+                module.__spec__ = self._spec
+                module.__loader__ = self._loader
+        class _NattenAliasFinder(_MetaPathFinder):
+            def find_spec(self, fullname, path=None, target=None):
+                if not fullname.startswith("natten."):
+                    return None
+                real_name = _real_root + fullname[len("natten") :]
+                try:
+                    module = _importlib.import_module(real_name)
+                except ImportError:
+                    return None
+                return _spec_from_loader(fullname, _NattenAliasLoader(module))
+        # Must precede PathFinder, which would otherwise re-execute
+        # submodules reachable through a real parent package's __path__.
+        _sys.meta_path.insert(0, _NattenAliasFinder())
+__all__ = [
+    "__version__",
+    "_private_for_testing",
+    "NeighborhoodAttention1D",
+    "NeighborhoodAttention2D",
+    "NeighborhoodAttention3D",
+    "are_deterministic_algorithms_enabled",
+    "use_deterministic_algorithms",
+    "use_kv_parallelism_in_fused_na",
+    "is_kv_parallelism_in_fused_na_enabled",
+    "set_memory_usage_preference",
+    "get_memory_usage_preference",
+    "is_memory_usage_default",
+    "is_memory_usage_strict",
+    "is_memory_usage_unrestricted",
+    "is_flex_compile_allowed",
+    "is_flex_compile_backprop_allowed",
+    "allow_flex_compile",
+    "allow_flex_compile_backprop",
+    "disable_flex_compile",
+    "disable_flex_compile_backprop",
+    "get_bwd_configs_for_cutlass_fmha",
+    "get_bwd_configs_for_cutlass_fna",
+    "get_configs_for_cutlass_fmha",
+    "get_configs_for_cutlass_fna",
+    "get_configs_for_cutlass_hopper_fmha",
+    "get_bwd_configs_for_cutlass_hopper_fmha",
+    "get_configs_for_cutlass_hopper_fna",
+    "get_bwd_configs_for_cutlass_hopper_fna",
+    "get_bwd_configs_for_cutlass_blackwell_fmha",
+    "get_bwd_configs_for_cutlass_blackwell_fna",
+    "get_configs_for_cutlass_blackwell_fmha",
+    "get_configs_for_cutlass_blackwell_fna",
+    "get_configs_for_flex_fmha",
+    "get_configs_for_flex_fna",
+    "HAS_LIBNATTEN",
+    "na1d",
+    "na2d",
+    "na3d",
+    "attention",
+    "merge_attentions",
+]

build/torch210-cu126-aarch64-linux/_environment.py ADDED Viewed

	@@ -0,0 +1,59 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+from ._libnatten import HAS_LIBNATTEN  # noqa: F401
+from .utils.environment import (
+    _IS_CUDA_AVAILABLE,
+    _IS_TORCH_COMPILE_SUPPORTED,
+    _TORCH_VERSION,
+    parse_env_flag,
+    parse_env_int,
+    parse_env_str,
+)
+# Default tokperm implementation; choices:
+# NATTEN_TOKPERM_DEFAULT_IMPL="cutlass"
+# NATTEN_TOKPERM_DEFAULT_IMPL="torch"
+USE_TORCH_IMPL_DEFAULT = (
+    parse_env_str("NATTEN_TOKPERM_DEFAULT_IMPL", "cutlass") == "torch"
+)
+# Unit tests
+_RUN_EXTENDED_TESTS = parse_env_flag("NATTEN_RUN_EXTENDED_TESTS", False)
+_RUN_FLEX_TESTS = parse_env_flag("NATTEN_RUN_FLEX_TESTS", True)
+_NUM_RAND_SWEEP_TESTS = parse_env_int("NATTEN_RAND_SWEEP_TESTS", 1000)
+# Profiler
+DISABLE_TQDM = parse_env_flag("NATTEN_DISABLE_TQDM", False)
+__all__ = [
+    "HAS_LIBNATTEN",
+    "_IS_CUDA_AVAILABLE",
+    "_IS_TORCH_COMPILE_SUPPORTED",
+    "DISABLE_TQDM",
+    "_RUN_FLEX_TESTS",
+    "_RUN_FLEX_TESTS",
+    "_NUM_RAND_SWEEP_TESTS",
+    "_TORCH_VERSION",
+]

build/torch210-cu126-aarch64-linux/_libnatten/__init__.py ADDED Viewed

	@@ -0,0 +1,109 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+# kernel-builder port: libnatten is always compiled into this kernel; the
+# upstream stub fallback path is not needed.
+import torch  # noqa: F401
+HAS_LIBNATTEN = True
+from .torch_wrappers import (
+    blackwell_fmha_backward,
+    blackwell_fmha_forward,
+    blackwell_na1d_backward,
+    blackwell_na1d_forward,
+    blackwell_na2d_backward,
+    blackwell_na2d_forward,
+    blackwell_na3d_backward,
+    blackwell_na3d_forward,
+    compute_delta,
+    fmha_backward,
+    fmha_forward,
+    hopper_fmha_backward,
+    hopper_fmha_forward,
+    hopper_na1d_backward,
+    hopper_na1d_forward,
+    hopper_na2d_backward,
+    hopper_na2d_forward,
+    hopper_na3d_backward,
+    hopper_na3d_forward,
+    na1d_backward,
+    na1d_forward,
+    na2d_backward,
+    na2d_forward,
+    na3d_backward,
+    na3d_forward,
+    reference_na1d_backward,
+    reference_na1d_forward,
+    reference_na2d_backward,
+    reference_na2d_forward,
+    reference_na3d_backward,
+    reference_na3d_forward,
+    token_permute_1d,
+    token_permute_2d,
+    token_permute_3d,
+    token_unpermute_1d,
+    token_unpermute_2d,
+    token_unpermute_3d,
+)
+__all__ = [
+    "HAS_LIBNATTEN",
+    "blackwell_fmha_backward",
+    "blackwell_fmha_forward",
+    "blackwell_na1d_backward",
+    "blackwell_na1d_forward",
+    "blackwell_na2d_backward",
+    "blackwell_na2d_forward",
+    "blackwell_na3d_backward",
+    "blackwell_na3d_forward",
+    "compute_delta",
+    "fmha_backward",
+    "fmha_forward",
+    "hopper_fmha_backward",
+    "hopper_fmha_forward",
+    "hopper_na1d_backward",
+    "hopper_na1d_forward",
+    "hopper_na2d_backward",
+    "hopper_na2d_forward",
+    "hopper_na3d_backward",
+    "hopper_na3d_forward",
+    "na1d_backward",
+    "na1d_forward",
+    "na2d_backward",
+    "na2d_forward",
+    "na3d_backward",
+    "na3d_forward",
+    "reference_na1d_backward",
+    "reference_na1d_forward",
+    "reference_na2d_backward",
+    "reference_na2d_forward",
+    "reference_na3d_backward",
+    "reference_na3d_forward",
+    "token_permute_1d",
+    "token_permute_2d",
+    "token_permute_3d",
+    "token_unpermute_1d",
+    "token_unpermute_2d",
+    "token_unpermute_3d",
+]

build/torch210-cu126-aarch64-linux/_libnatten/torch_wrappers.py ADDED Viewed

	@@ -0,0 +1,1006 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+# kernel-builder port of upstream `natten/_libnatten/torch_wrappers.py`.
+#
+# Upstream registers Python `torch.library.custom_op`s that allocate outputs
+# and call into the pybind11 `libnatten` extension. In this port the ops are
+# registered in C++ (`torch-ext/torch_binding.cpp`) as out-variant ops under
+# the build-time namespace exposed through `.._ops`. The functions here keep
+# the exact upstream calling conventions (allocate outputs, handle kv-split
+# defaults, varlen zero-init) and call the C++ ops, and each C++ op gets a
+# fake (meta) registration so the whole surface stays torch.compile-safe.
+#
+# Schema conventions of the C++ ops:
+#   - `kernel_size`/`stride`/`dilation`/tile shapes are `int[]`.
+#   - Multi-dimensional causal masks are passed as `int[]` (0/1) because
+#     boolean arrays are less uniformly supported in op schemas.
+#   - `scale` is a `float`.
+import math
+from typing import Optional, Sequence, Tuple
+import torch
+from torch import Tensor
+from .._ops import add_op_namespace_prefix, ops
+from ..utils.tuples import ceil_div_tuple, mul_tuple
+register_fake = torch.library.register_fake
+def maybe_contiguous(x):
+    return x.contiguous()
+def _ints(v: Sequence) -> list:
+    return [int(x) for x in v]
+################################################################################
+############################ Fake (meta) registration ##########################
+################################################################################
+# All C++ ops are out-variant: they only mutate output arguments and return
+# nothing, so their fake impls are no-ops. Shape inference happens in the
+# Python wrappers below, which allocate the outputs.
+def _register_noop_fake(op_name: str) -> None:
+    def _fake(*args, **kwargs) -> None:
+        return None
+    register_fake(add_op_namespace_prefix(op_name))(_fake)
+for _na_dim in (1, 2, 3):
+    for _prefix in ("", "hopper_", "blackwell_", "reference_"):
+        _register_noop_fake(f"{_prefix}na{_na_dim}d_forward")
+        _register_noop_fake(f"{_prefix}na{_na_dim}d_backward")
+    _register_noop_fake(f"token_permute_{_na_dim}d")
+    _register_noop_fake(f"token_unpermute_{_na_dim}d")
+for _prefix in ("", "hopper_", "blackwell_"):
+    _register_noop_fake(f"{_prefix}fmha_forward")
+    _register_noop_fake(f"{_prefix}fmha_backward")
+_register_noop_fake("compute_delta")
+################################################################################
+################################### FMHA ops ###################################
+################################################################################
+def blackwell_fmha_forward(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    is_causal: bool,
+    scale: float,
+    q_tile_size: int,
+    kv_tile_size: int,
+    run_persistent_kernel: bool,
+    cumulative_seqlen_Q: Optional[Tensor],
+    cumulative_seqlen_KV: Optional[Tensor],
+    max_seqlen_Q: int,
+    max_seqlen_KV: int,
+) -> Tuple[Tensor, Tensor]:
+    query, key, value = [maybe_contiguous(x) for x in (query, key, value)]
+    output_shape = [s for s in query.shape[:-1]] + [value.shape[-1]]
+    # NOTE: always zero-init outputs when doing varlen for safety
+    is_varlen = cumulative_seqlen_Q is not None
+    init_fn = torch.zeros if is_varlen else torch.empty
+    output = init_fn(output_shape, device=query.device, dtype=query.dtype)
+    logsumexp = init_fn(query.shape[:-1], dtype=torch.float32, device=query.device)
+    # Skip kernel launch when all sequences are empty
+    if is_varlen and max_seqlen_Q == 0 and max_seqlen_KV == 0:
+        return output, logsumexp
+    ops.blackwell_fmha_forward(
+        output,
+        query,
+        key,
+        value,
+        logsumexp,
+        bool(is_causal),
+        float(scale),
+        int(q_tile_size),
+        int(kv_tile_size),
+        bool(run_persistent_kernel),
+        cumulative_seqlen_Q,
+        cumulative_seqlen_KV,
+        int(max_seqlen_Q),
+        int(max_seqlen_KV),
+    )
+    return output, logsumexp
+def blackwell_fmha_backward(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    output: Tensor,
+    d_output: Tensor,
+    logsumexp: Tensor,
+    is_causal: bool,
+    scale: float,
+    q_tile_size: int,
+    kv_tile_size: int,
+    cumulative_seqlen_Q: Optional[Tensor],
+    cumulative_seqlen_KV: Optional[Tensor],
+    max_seqlen_Q: int,
+    max_seqlen_KV: int,
+    deterministic: bool,
+) -> Tuple[Tensor, Tensor, Tensor]:
+    query, key, value = [maybe_contiguous(x) for x in (query, key, value)]
+    output, d_output, logsumexp = [
+        maybe_contiguous(x) for x in (output, d_output, logsumexp)
+    ]
+    # NOTE: always zero-init outputs when doing varlen for safety
+    is_varlen = cumulative_seqlen_Q is not None
+    init_fn = torch.zeros_like if is_varlen else torch.empty_like
+    d_query = init_fn(query)
+    d_key = init_fn(key)
+    d_value = init_fn(value)
+    # Skip kernel launch when all sequences are empty
+    if is_varlen and max_seqlen_Q == 0 and max_seqlen_KV == 0:
+        return d_query, d_key, d_value
+    ops.blackwell_fmha_backward(
+        d_query,
+        d_key,
+        d_value,
+        query,
+        key,
+        value,
+        output,
+        d_output,
+        logsumexp,
+        bool(is_causal),
+        float(scale),
+        int(q_tile_size),
+        int(kv_tile_size),
+        cumulative_seqlen_Q,
+        cumulative_seqlen_KV,
+        int(max_seqlen_Q),
+        int(max_seqlen_KV),
+        bool(deterministic),
+    )
+    return d_query, d_key, d_value
+def hopper_fmha_forward(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    is_causal: bool,
+    scale: float,
+    q_tile_size: int,
+    kv_tile_size: int,
+    kernel_schedule_int: int,
+    cumulative_seqlen_Q: Optional[Tensor],
+    cumulative_seqlen_KV: Optional[Tensor],
+    max_seqlen_Q: int,
+    max_seqlen_KV: int,
+) -> Tuple[Tensor, Tensor]:
+    query, key, value = [maybe_contiguous(x) for x in (query, key, value)]
+    output_shape = [s for s in query.shape[:-1]] + [value.shape[-1]]
+    # NOTE: always zero-init outputs when doing varlen for safety
+    is_varlen = cumulative_seqlen_Q is not None
+    init_fn = torch.zeros if is_varlen else torch.empty
+    output = init_fn(output_shape, device=query.device, dtype=query.dtype)
+    logsumexp = init_fn(query.shape[:-1], dtype=torch.float32, device=query.device)
+    # Skip kernel launch when all sequences are empty
+    if is_varlen and max_seqlen_Q == 0 and max_seqlen_KV == 0:
+        return output, logsumexp
+    ops.hopper_fmha_forward(
+        output,
+        query,
+        key,
+        value,
+        logsumexp,
+        bool(is_causal),
+        float(scale),
+        int(q_tile_size),
+        int(kv_tile_size),
+        int(kernel_schedule_int),
+        cumulative_seqlen_Q,
+        cumulative_seqlen_KV,
+        int(max_seqlen_Q),
+        int(max_seqlen_KV),
+    )
+    return output, logsumexp
+def hopper_fmha_backward(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    output: Tensor,
+    d_output: Tensor,
+    logsumexp: Tensor,
+    is_causal: bool,
+    scale: float,
+    q_tile_size: int,
+    kv_tile_size: int,
+    cumulative_seqlen_Q: Optional[Tensor],
+    cumulative_seqlen_KV: Optional[Tensor],
+    max_seqlen_Q: int,
+    max_seqlen_KV: int,
+) -> Tuple[Tensor, Tensor, Tensor]:
+    query, key, value = [maybe_contiguous(x) for x in (query, key, value)]
+    output, d_output, logsumexp = [
+        maybe_contiguous(x) for x in (output, d_output, logsumexp)
+    ]
+    # NOTE: always zero-init outputs when doing varlen for safety
+    is_varlen = cumulative_seqlen_Q is not None
+    init_fn = torch.zeros_like if is_varlen else torch.empty_like
+    d_query = init_fn(query)
+    d_key = init_fn(key)
+    d_value = init_fn(value)
+    # Skip kernel launch when all sequences are empty
+    if is_varlen and max_seqlen_Q == 0 and max_seqlen_KV == 0:
+        return d_query, d_key, d_value
+    ops.hopper_fmha_backward(
+        d_query,
+        d_key,
+        d_value,
+        query,
+        key,
+        value,
+        output,
+        d_output,
+        logsumexp,
+        bool(is_causal),
+        float(scale),
+        int(q_tile_size),
+        int(kv_tile_size),
+        cumulative_seqlen_Q,
+        cumulative_seqlen_KV,
+        int(max_seqlen_Q),
+        int(max_seqlen_KV),
+    )
+    return d_query, d_key, d_value
+def fmha_forward(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    is_causal: bool,
+    scale: float,
+    q_tile_size: int,
+    kv_tile_size: int,
+    cumulative_seqlen_Q: Optional[Tensor],
+    cumulative_seqlen_KV: Optional[Tensor],
+    max_seqlen_Q: int,
+    max_seqlen_KV: int,
+) -> Tuple[Tensor, Tensor]:
+    query, key, value = [maybe_contiguous(x) for x in (query, key, value)]
+    output_shape = [s for s in query.shape[:-1]] + [value.shape[-1]]
+    # NOTE: always zero-init outputs when doing varlen for safety
+    is_varlen = cumulative_seqlen_Q is not None
+    init_fn = torch.zeros if is_varlen else torch.empty
+    output = init_fn(output_shape, device=query.device, dtype=query.dtype)
+    logsumexp = init_fn(query.shape[:-1], dtype=torch.float32, device=query.device)
+    # Skip kernel launch when all sequences are empty
+    if is_varlen and max_seqlen_Q == 0 and max_seqlen_KV == 0:
+        return output, logsumexp
+    ops.fmha_forward(
+        output,
+        query,
+        key,
+        value,
+        logsumexp,
+        bool(is_causal),
+        float(scale),
+        int(q_tile_size),
+        int(kv_tile_size),
+        cumulative_seqlen_Q,
+        cumulative_seqlen_KV,
+        int(max_seqlen_Q),
+        int(max_seqlen_KV),
+    )
+    return output, logsumexp
+def fmha_backward(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    output: Tensor,
+    d_output: Tensor,
+    logsumexp: Tensor,
+    is_causal: bool,
+    scale: float,
+    q_tile_size: int,
+    kv_tile_size: int,
+    num_kv_splits: Optional[int],
+    compute_delta_with_pt: bool,
+    cumulative_seqlen_Q: Optional[Tensor],
+    cumulative_seqlen_KV: Optional[Tensor],
+    max_seqlen_Q: int,
+    max_seqlen_KV: int,
+    deterministic: bool,
+) -> Tuple[Tensor, Tensor, Tensor]:
+    from ..backends.configs.cutlass.backward_knobs import check_fmha_kv_splits
+    query, key, value = [maybe_contiguous(x) for x in (query, key, value)]
+    output, d_output, logsumexp = [
+        maybe_contiguous(x) for x in (output, d_output, logsumexp)
+    ]
+    # NOTE: always zero-init outputs when doing varlen for safety
+    is_varlen = cumulative_seqlen_Q is not None
+    init_fn = torch.zeros_like if is_varlen else torch.empty_like
+    d_query = init_fn(query)
+    d_key = init_fn(key)
+    d_value = init_fn(value)
+    # Skip kernel launch when all sequences are empty
+    if is_varlen and max_seqlen_Q == 0 and max_seqlen_KV == 0:
+        return d_query, d_key, d_value
+    if deterministic:
+        # Torch reduction seems to have slight reproducibility issues, even with determinism on
+        compute_delta_with_pt = False
+        # TODO: this is the only way to get determinism in this kernel, but it's very slow
+        num_kv_splits = 1
+    else:
+        # Compute default kv_splits if not specified
+        # max_seqlen must be at least 2 to satisfy static checks that are just too complicated to
+        # relax at this point. Kernel launch will be skipped if max_seqlen is 0 anyway. Prior checks
+        # should prevent negative max seqlens.
+        max_seqlen = max(2, max_seqlen_KV) if is_varlen else None
+        num_kv_splits = check_fmha_kv_splits(
+            kv_splits=num_kv_splits,
+            input_tensor=key,
+            kv_tile_size=kv_tile_size,
+            deterministic=deterministic,
+            max_seqlen=max_seqlen,
+        )
+    ops.fmha_backward(
+        d_query,
+        d_key,
+        d_value,
+        query,
+        key,
+        value,
+        output,
+        d_output,
+        logsumexp,
+        bool(is_causal),
+        float(scale),
+        int(q_tile_size),
+        int(kv_tile_size),
+        int(num_kv_splits),
+        bool(compute_delta_with_pt),
+        cumulative_seqlen_Q,
+        cumulative_seqlen_KV,
+        int(max_seqlen_Q),
+        int(max_seqlen_KV),
+    )
+    return d_query, d_key, d_value
+################################################################################
+################################### FNA ops  ###################################
+################################################################################
+def make_blackwell_fna_ops(na_dim):
+    fwd_op = getattr(ops, f"blackwell_na{na_dim}d_forward")
+    bwd_op = getattr(ops, f"blackwell_na{na_dim}d_backward")
+    def blackwell_fna_forward(
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        kernel_size,
+        stride,
+        dilation,
+        is_causal,
+        scale: float,
+        q_shape,
+        kv_shape,
+        qkv_shape,
+        q_tile_shape,
+        kv_tile_shape,
+        run_persistent_kernel: bool,
+    ) -> Tuple[Tensor, Tensor]:
+        query, key, value = [maybe_contiguous(x) for x in (query, key, value)]
+        output_shape = [s for s in query.shape[:-1]] + [value.shape[-1]]
+        output = torch.empty(output_shape, device=query.device, dtype=query.dtype)
+        logsumexp = torch.empty(
+            query.shape[:-1], dtype=torch.float32, device=query.device
+        )
+        fwd_op(
+            output,
+            query,
+            key,
+            value,
+            logsumexp,
+            _ints(kernel_size),
+            _ints(stride),
+            _ints(dilation),
+            _ints(is_causal),
+            float(scale),
+            _ints(q_shape),
+            _ints(kv_shape),
+            _ints(qkv_shape),
+            _ints(q_tile_shape),
+            _ints(kv_tile_shape),
+            bool(run_persistent_kernel),
+        )
+        return output, logsumexp
+    def blackwell_fna_backward(
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        output: Tensor,
+        d_output: Tensor,
+        logsumexp: Tensor,
+        kernel_size,
+        stride,
+        dilation,
+        is_causal,
+        scale: float,
+        q_shape,
+        kv_shape,
+        qkv_shape,
+        q_tile_shape,
+        kv_tile_shape,
+    ) -> Tuple[Tensor, Tensor, Tensor]:
+        query, key, value = [maybe_contiguous(x) for x in (query, key, value)]
+        output, d_output, logsumexp = [
+            maybe_contiguous(x) for x in (output, d_output, logsumexp)
+        ]
+        d_query = torch.empty_like(query)
+        d_key = torch.empty_like(key)
+        d_value = torch.empty_like(value)
+        bwd_op(
+            d_query,
+            d_key,
+            d_value,
+            query,
+            key,
+            value,
+            output,
+            d_output,
+            logsumexp,
+            _ints(kernel_size),
+            _ints(stride),
+            _ints(dilation),
+            _ints(is_causal),
+            float(scale),
+            _ints(q_shape),
+            _ints(kv_shape),
+            _ints(qkv_shape),
+            _ints(q_tile_shape),
+            _ints(kv_tile_shape),
+        )
+        return d_query, d_key, d_value
+    return blackwell_fna_forward, blackwell_fna_backward
+def make_hopper_fna_ops(na_dim):
+    fwd_op = getattr(ops, f"hopper_na{na_dim}d_forward")
+    bwd_op = getattr(ops, f"hopper_na{na_dim}d_backward")
+    def hopper_fna_forward(
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        kernel_size,
+        stride,
+        dilation,
+        is_causal,
+        scale: float,
+        q_shape,
+        kv_shape,
+        qkv_shape,
+        q_tile_shape,
+        kv_tile_shape,
+        kernel_schedule_int: int,
+    ) -> Tuple[Tensor, Tensor]:
+        query, key, value = [maybe_contiguous(x) for x in (query, key, value)]
+        output_shape = [s for s in query.shape[:-1]] + [value.shape[-1]]
+        output = torch.empty(output_shape, device=query.device, dtype=query.dtype)
+        logsumexp = torch.empty(
+            query.shape[:-1], dtype=torch.float32, device=query.device
+        )
+        fwd_op(
+            output,
+            query,
+            key,
+            value,
+            logsumexp,
+            _ints(kernel_size),
+            _ints(stride),
+            _ints(dilation),
+            _ints(is_causal),
+            float(scale),
+            _ints(q_shape),
+            _ints(kv_shape),
+            _ints(qkv_shape),
+            _ints(q_tile_shape),
+            _ints(kv_tile_shape),
+            int(kernel_schedule_int),
+        )
+        return output, logsumexp
+    def hopper_fna_backward(
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        output: Tensor,
+        d_output: Tensor,
+        logsumexp: Tensor,
+        kernel_size,
+        stride,
+        dilation,
+        is_causal,
+        scale: float,
+        q_shape,
+        kv_shape,
+        qkv_shape,
+        q_tile_shape,
+        kv_tile_shape,
+    ) -> Tuple[Tensor, Tensor, Tensor]:
+        query, key, value = [maybe_contiguous(x) for x in (query, key, value)]
+        output, d_output, logsumexp = [
+            maybe_contiguous(x) for x in (output, d_output, logsumexp)
+        ]
+        d_query = torch.empty_like(query)
+        d_key = torch.empty_like(key)
+        d_value = torch.empty_like(value)
+        bwd_op(
+            d_query,
+            d_key,
+            d_value,
+            query,
+            key,
+            value,
+            output,
+            d_output,
+            logsumexp,
+            _ints(kernel_size),
+            _ints(stride),
+            _ints(dilation),
+            _ints(is_causal),
+            float(scale),
+            _ints(q_shape),
+            _ints(kv_shape),
+            _ints(qkv_shape),
+            _ints(q_tile_shape),
+            _ints(kv_tile_shape),
+        )
+        return d_query, d_key, d_value
+    return hopper_fna_forward, hopper_fna_backward
+def make_fna_ops(na_dim):
+    fwd_op = getattr(ops, f"na{na_dim}d_forward")
+    bwd_op = getattr(ops, f"na{na_dim}d_backward")
+    def fna_forward(
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        kernel_size,
+        stride,
+        dilation,
+        is_causal,
+        scale: float,
+        q_tile_shape,
+        kv_tile_shape,
+    ) -> Tuple[Tensor, Tensor]:
+        query, key, value = [maybe_contiguous(x) for x in (query, key, value)]
+        output_shape = [s for s in query.shape[:-1]] + [value.shape[-1]]
+        output = torch.empty(output_shape, device=query.device, dtype=query.dtype)
+        logsumexp = torch.empty(
+            query.shape[:-1], dtype=torch.float32, device=query.device
+        )
+        fwd_op(
+            output,
+            query,
+            key,
+            value,
+            logsumexp,
+            _ints(kernel_size),
+            _ints(stride),
+            _ints(dilation),
+            _ints(is_causal),
+            float(scale),
+            _ints(q_tile_shape),
+            _ints(kv_tile_shape),
+        )
+        return output, logsumexp
+    def fna_backward(
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        output: Tensor,
+        d_output: Tensor,
+        logsumexp: Tensor,
+        kernel_size,
+        stride,
+        dilation,
+        is_causal,
+        scale: float,
+        q_tile_shape,
+        kv_tile_shape,
+        num_kv_splits,
+        compute_delta_with_pt: bool,
+        deterministic: bool,
+    ) -> Tuple[Tensor, Tensor, Tensor]:
+        from ..backends.configs.cutlass.backward_knobs import check_fna_kv_splits
+        query, key, value = [maybe_contiguous(x) for x in (query, key, value)]
+        output, d_output, logsumexp = [
+            maybe_contiguous(x) for x in (output, d_output, logsumexp)
+        ]
+        d_query = torch.empty_like(query)
+        d_key = torch.empty_like(key)
+        d_value = torch.empty_like(value)
+        if deterministic:
+            # Torch reduction seems to have slight reproducibility issues, even with determinism on
+            compute_delta_with_pt = False
+            # TODO: this is the only way to get determinism in this kernel, but it's very slow
+            num_kv_splits = tuple(1 for _ in range(na_dim))
+        else:
+            # Compute default kv_splits if not specified
+            num_kv_splits = check_fna_kv_splits(
+                kv_splits=tuple(num_kv_splits) if num_kv_splits is not None else None,
+                input_tensor=key,
+                kv_tile_shape=tuple(kv_tile_shape),
+                deterministic=deterministic,
+                dilation=tuple(dilation),
+            )
+        bwd_op(
+            d_query,
+            d_key,
+            d_value,
+            query,
+            key,
+            value,
+            output,
+            d_output,
+            logsumexp,
+            _ints(kernel_size),
+            _ints(stride),
+            _ints(dilation),
+            _ints(is_causal),
+            float(scale),
+            _ints(q_tile_shape),
+            _ints(kv_tile_shape),
+            _ints(num_kv_splits),
+            bool(compute_delta_with_pt),
+        )
+        return d_query, d_key, d_value
+    return fna_forward, fna_backward
+def make_reference_fna_ops(na_dim):
+    fwd_op = getattr(ops, f"reference_na{na_dim}d_forward")
+    bwd_op = getattr(ops, f"reference_na{na_dim}d_backward")
+    def reference_fna_forward(
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        kernel_size,
+        stride,
+        dilation,
+        is_causal,
+        scale: float,
+        qkv_shape,
+        num_extra_kv: int,
+    ) -> Tuple[Tensor, Tensor]:
+        query, key, value = [maybe_contiguous(x) for x in (query, key, value)]
+        output_shape = [s for s in query.shape[:-1]] + [value.shape[-1]]
+        output = torch.empty(output_shape, device=query.device, dtype=query.dtype)
+        logsumexp = torch.empty(
+            query.shape[:-1], dtype=torch.float32, device=query.device
+        )
+        fwd_op(
+            output,
+            query,
+            key,
+            value,
+            logsumexp,
+            _ints(kernel_size),
+            _ints(stride),
+            _ints(dilation),
+            _ints(is_causal),
+            float(scale),
+            _ints(qkv_shape),
+            int(num_extra_kv),
+        )
+        return output, logsumexp
+    def reference_fna_backward(
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        output: Tensor,
+        d_output: Tensor,
+        logsumexp: Tensor,
+        kernel_size,
+        stride,
+        dilation,
+        is_causal,
+        scale: float,
+        qkv_shape,
+        num_extra_kv: int,
+    ) -> Tuple[Tensor, Tensor, Tensor]:
+        query, key, value = [maybe_contiguous(x) for x in (query, key, value)]
+        output, d_output, logsumexp = [
+            maybe_contiguous(x) for x in (output, d_output, logsumexp)
+        ]
+        d_query = torch.empty_like(query)
+        d_key = torch.empty_like(key)
+        d_value = torch.empty_like(value)
+        bwd_op(
+            d_query,
+            d_key,
+            d_value,
+            query,
+            key,
+            value,
+            output,
+            d_output,
+            logsumexp,
+            _ints(kernel_size),
+            _ints(stride),
+            _ints(dilation),
+            _ints(is_causal),
+            float(scale),
+            _ints(qkv_shape),
+            int(num_extra_kv),
+        )
+        return d_query, d_key, d_value
+    return reference_fna_forward, reference_fna_backward
+################################################################################
+################################# TokPerm ops  #################################
+################################################################################
+def make_token_permute_ops(na_dim):
+    permute_op = getattr(ops, f"token_permute_{na_dim}d")
+    unpermute_op = getattr(ops, f"token_unpermute_{na_dim}d")
+    def token_permute(
+        input_tensor: Tensor,
+        tile_shape,
+        dilation,
+        flip_tiled_dims: bool,
+    ) -> Tensor:
+        input_tensor = maybe_contiguous(input_tensor)
+        token_layout = tuple(x for x in input_tensor.shape[1 : na_dim + 1])
+        token_layout_padded = mul_tuple(
+            mul_tuple(
+                ceil_div_tuple(ceil_div_tuple(token_layout, tile_shape), dilation),
+                dilation,
+            ),
+            tile_shape,
+        )
+        output_shape = [
+            input_tensor.shape[0],
+            math.prod(token_layout_padded),
+            input_tensor.shape[-2],
+            input_tensor.shape[-1],
+        ]
+        output = torch.empty(
+            output_shape, device=input_tensor.device, dtype=input_tensor.dtype
+        )
+        permute_op(
+            output,
+            input_tensor,
+            _ints(tile_shape),
+            _ints(dilation),
+            bool(flip_tiled_dims),
+        )
+        # Fold dilation in batch dimension so that attention is correct.
+        output = output.reshape(
+            input_tensor.shape[0] * math.prod(dilation),
+            -1,
+            input_tensor.shape[-2],
+            input_tensor.shape[-1],
+        )
+        return output
+    def token_unpermute(
+        input_tensor: Tensor,
+        token_layout_shape,
+        tile_shape,
+        dilation,
+        flip_tiled_dims: bool,
+    ) -> Tensor:
+        input_tensor = maybe_contiguous(input_tensor)
+        # Unfold dilation in batch dimension
+        num_dilation_groups = math.prod(dilation)
+        assert input_tensor.shape[0] % num_dilation_groups == 0
+        input_tensor = input_tensor.reshape(
+            input_tensor.shape[0] // num_dilation_groups,
+            -1,
+            input_tensor.shape[-2],
+            input_tensor.shape[-1],
+        )
+        output_shape = [
+            input_tensor.shape[0],
+            *token_layout_shape,
+            input_tensor.shape[-2],
+            input_tensor.shape[-1],
+        ]
+        output = torch.empty(
+            output_shape, device=input_tensor.device, dtype=input_tensor.dtype
+        )
+        unpermute_op(
+            output,
+            input_tensor,
+            _ints(tile_shape),
+            _ints(dilation),
+            bool(flip_tiled_dims),
+        )
+        return output
+    return token_permute, token_unpermute
+(blackwell_na1d_forward, blackwell_na1d_backward) = make_blackwell_fna_ops(1)
+(blackwell_na2d_forward, blackwell_na2d_backward) = make_blackwell_fna_ops(2)
+(blackwell_na3d_forward, blackwell_na3d_backward) = make_blackwell_fna_ops(3)
+(hopper_na1d_forward, hopper_na1d_backward) = make_hopper_fna_ops(1)
+(hopper_na2d_forward, hopper_na2d_backward) = make_hopper_fna_ops(2)
+(hopper_na3d_forward, hopper_na3d_backward) = make_hopper_fna_ops(3)
+(na1d_forward, na1d_backward) = make_fna_ops(1)
+(na2d_forward, na2d_backward) = make_fna_ops(2)
+(na3d_forward, na3d_backward) = make_fna_ops(3)
+(reference_na1d_forward, reference_na1d_backward) = make_reference_fna_ops(1)
+(reference_na2d_forward, reference_na2d_backward) = make_reference_fna_ops(2)
+(reference_na3d_forward, reference_na3d_backward) = make_reference_fna_ops(3)
+(token_permute_1d, token_unpermute_1d) = make_token_permute_ops(1)
+(token_permute_2d, token_unpermute_2d) = make_token_permute_ops(2)
+(token_permute_3d, token_unpermute_3d) = make_token_permute_ops(3)
+# This is only used in unit tests, and not even auto-diffable
+def compute_delta(out: Tensor, d_out: Tensor, delta: Tensor) -> None:
+    ops.compute_delta(out, d_out, delta)
+__all__ = [
+    "blackwell_fmha_backward",
+    "blackwell_fmha_forward",
+    "blackwell_na1d_backward",
+    "blackwell_na1d_forward",
+    "blackwell_na2d_backward",
+    "blackwell_na2d_forward",
+    "blackwell_na3d_backward",
+    "blackwell_na3d_forward",
+    "compute_delta",
+    "fmha_backward",
+    "fmha_forward",
+    "hopper_fmha_backward",
+    "hopper_fmha_forward",
+    "hopper_na1d_backward",
+    "hopper_na1d_forward",
+    "hopper_na2d_backward",
+    "hopper_na2d_forward",
+    "hopper_na3d_backward",
+    "hopper_na3d_forward",
+    "na1d_backward",
+    "na1d_forward",
+    "na2d_backward",
+    "na2d_forward",
+    "na3d_backward",
+    "na3d_forward",
+    "reference_na1d_backward",
+    "reference_na1d_forward",
+    "reference_na2d_backward",
+    "reference_na2d_forward",
+    "reference_na3d_backward",
+    "reference_na3d_forward",
+    "token_permute_1d",
+    "token_permute_2d",
+    "token_permute_3d",
+    "token_unpermute_1d",
+    "token_unpermute_2d",
+    "token_unpermute_3d",
+]

build/torch210-cu126-aarch64-linux/_natten_cuda_5fe9dbb.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:02f9f3cca0fe961f460ce848017cc7a391aa74e34ca29e5f6431d1138c433600
+size 100978320

build/torch210-cu126-aarch64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _natten_cuda_5fe9dbb
+ops = torch.ops._natten_cuda_5fe9dbb
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_natten_cuda_5fe9dbb::{op_name}"

build/torch210-cu126-aarch64-linux/_private_for_testing.py ADDED Viewed

	@@ -0,0 +1,33 @@

+from . import (
+    _environment,
+    _libnatten,
+    _ops,
+    _types,
+    backends,
+    context,
+    functional,
+    token_permute,
+)
+from .backends import reference
+from .backends.configs import cutlass, cutlass_blackwell, cutlass_hopper
+from .utils import checks, dtype, log, testing, varlen
+__all__ = [
+    "_environment",
+    "_libnatten",
+    "_ops",
+    "_types",
+    "backends",
+    "checks",
+    "context",
+    "cutlass",
+    "cutlass_blackwell",
+    "cutlass_hopper",
+    "dtype",
+    "functional",
+    "log",
+    "reference",
+    "testing",
+    "token_permute",
+    "varlen",
+]

build/torch210-cu126-aarch64-linux/_types.py ADDED Viewed

	@@ -0,0 +1,85 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+from enum import Enum
+from typing import Tuple, Union
+NoneType = type(None)
+Dimension1DType = Tuple[int]
+Dimension2DType = Tuple[int, int]
+Dimension3DType = Tuple[int, int, int]
+CausalArg1DType = Tuple[bool]
+CausalArg2DType = Tuple[bool, bool]
+CausalArg3DType = Tuple[bool, bool, bool]
+# NOTE: switch to | when < 3.10 support is dropped
+Dimension1DTypeOrDed = Union[int, Dimension1DType]
+Dimension2DTypeOrDed = Union[int, Dimension2DType]
+Dimension3DTypeOrDed = Union[int, Dimension3DType]
+CausalArg1DTypeOrDed = Union[bool, CausalArg1DType]
+CausalArg2DTypeOrDed = Union[bool, CausalArg2DType]
+CausalArg3DTypeOrDed = Union[bool, CausalArg3DType]
+DimensionType = Union[Dimension1DType, Dimension2DType, Dimension3DType]
+CausalArgType = Union[CausalArg1DType, CausalArg2DType, CausalArg3DType]
+DimensionTypeOrDed = Union[int, DimensionType]
+CausalArgTypeOrDed = Union[bool, CausalArgType]
+# (query_tile_shape, kv_tile_shape)
+QKTileShapeType = Union[
+    Tuple[Dimension1DType, Dimension1DType],
+    Tuple[Dimension2DType, Dimension2DType],
+    Tuple[Dimension3DType, Dimension3DType],
+]
+# TODO: Only applies to Hopper FMHA/FNA for now -- extend to other applicable kernels
+class KernelSchedule(Enum):
+    NonPersistent = 0
+    WarpSpecializedCooperative = 1
+    WarpSpecializedPingpong = 2
+CutlassFnaForwardConfigType = QKTileShapeType
+CutlassFnaBackwardConfigType = QKTileShapeType
+CutlassBlackwellFnaForwardConfigType = QKTileShapeType
+CutlassBlackwellFnaBackwardConfigType = QKTileShapeType
+CutlassHopperFnaForwardConfigType = Tuple[QKTileShapeType, KernelSchedule]
+CutlassHopperFnaBackwardConfigType = QKTileShapeType
+FlexFnaForwardConfigType = QKTileShapeType
+# FMHA configs
+FmhaForwardConfigType = Tuple[int, int]
+CutlassFmhaForwardConfigType = FmhaForwardConfigType
+CutlassFmhaBackwardConfigType = FmhaForwardConfigType
+FlexFmhaForwardConfigType = FmhaForwardConfigType
+CutlassBlackwellFmhaForwardConfigType = FmhaForwardConfigType
+CutlassBlackwellFmhaBackwardConfigType = FmhaForwardConfigType
+CutlassHopperFmhaForwardConfigType = Tuple[FmhaForwardConfigType, KernelSchedule]
+CutlassHopperFmhaBackwardConfigType = FmhaForwardConfigType

build/torch210-cu126-aarch64-linux/attn_merge.py ADDED Viewed

	@@ -0,0 +1,292 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+import functools
+from typing import List, Tuple
+import torch
+from torch import Tensor
+from torch.amp import custom_bwd, custom_fwd
+from torch.autograd import Function
+amp_fwd = functools.partial(custom_fwd, device_type="cuda")
+amp_bwd = functools.partial(custom_bwd, device_type="cuda")
+from ._environment import _IS_TORCH_COMPILE_SUPPORTED
+def _maybe_torch_compile(*args, **kwargs):
+    def decorator(f):
+        if _IS_TORCH_COMPILE_SUPPORTED:
+            return torch.compile(f, *args, **kwargs)
+        return f
+    return decorator
+# TODO: if use cases for this grow, we might want to do a custom kernel
+def _merge_attentions_fn(
+    outputs: List[Tensor], lse_tensors: List[Tensor]
+) -> Tuple[Tensor, Tensor]:
+    assert len(outputs) >= 2, "Expected at least two tensors."
+    num_splits = len(outputs)
+    assert (
+        len(lse_tensors) == num_splits
+    ), "Expected number of outputs and LSE tensors to match."
+    assert all(
+        output.dim() == 4 and output.is_contiguous() for output in outputs
+    ), "Output tensors must be rank-4 tensors with (batch, seq, heads, dim) contiguous layout."
+    batch, seqlen, heads, dim = outputs[0].shape
+    assert all(
+        [x for x in output.shape] == [batch, seqlen, heads, dim] for output in outputs
+    ), "Output tensors must match in shape."
+    assert all(
+        lse.dim() == 3
+        and lse.is_contiguous()
+        and [x for x in lse.shape] == [batch, seqlen, heads]
+        for lse in lse_tensors
+    ), "LSE tensors must be rank-3 tensors with (batch, seq, heads) contiguous layout, and match in shape."
+    accum_type = torch.float32
+    output_type = outputs[0].dtype
+    lse_tensors = [lse.to(accum_type).unsqueeze(-1) for lse in lse_tensors]
+    outputs = [output.to(accum_type) for output in outputs]
+    # New approach based on https://github.com/zhuzilin/ring-flash-attention/pull/34
+    output = outputs[0] - torch.nn.functional.sigmoid(
+        lse_tensors[1] - lse_tensors[0]
+    ) * (outputs[0] - outputs[1])
+    logsumexp = lse_tensors[0] - torch.nn.functional.logsigmoid(
+        lse_tensors[0] - lse_tensors[1]
+    )
+    for i in range(2, num_splits):
+        output = output - torch.nn.functional.sigmoid(lse_tensors[i] - logsumexp) * (
+            output - outputs[i]
+        )
+        logsumexp = logsumexp - torch.nn.functional.logsigmoid(
+            logsumexp - lse_tensors[i]
+        )
+    output = output.to(output_type)
+    logsumexp = logsumexp.squeeze(-1)
+    assert logsumexp.dim() == 3
+    assert logsumexp.shape[0] == batch
+    assert logsumexp.shape[1] == seqlen
+    assert logsumexp.shape[2] == heads
+    return output, logsumexp
+@_maybe_torch_compile(fullgraph=True)
+def _merge_attentions_compile(
+    outputs: List[Tensor], lse_tensors: List[Tensor]
+) -> Tuple[Tensor, Tensor]:
+    return _merge_attentions_fn(outputs, lse_tensors)
+def _merge_attentions_op(
+    outputs: List[Tensor], lse_tensors: List[Tensor], torch_compile: bool = True
+) -> Tuple[Tensor, Tensor]:
+    if not torch_compile:
+        return _merge_attentions_fn(
+            [output.contiguous() for output in outputs],
+            [lse.contiguous() for lse in lse_tensors],
+        )
+    return _merge_attentions_compile(
+        [output.contiguous() for output in outputs],
+        [lse.contiguous() for lse in lse_tensors],
+    )
+class MergeAttentionsAutogradFn(Function):
+    @staticmethod
+    @amp_fwd
+    def forward(
+        ctx,
+        *args,
+    ) -> Tuple[Tensor, Tensor]:
+        assert len(args) >= 5, (
+            "Expected at least 5 args (two outputs, two lse tensors, 1 torch compile flag) "
+            + f"in attention merge, got {len(args)}."
+        )
+        assert (len(args) - 1) % 2 == 0, (
+            f"Expected pairs of outputs and lse tensors, got {len(args)-1} args "
+            + "(excluding torch compile flag)"
+        )
+        num_pairs = (len(args) - 1) // 2
+        assert num_pairs >= 2
+        torch_compile = args[-1]
+        outputs = args[:num_pairs]
+        lses = args[num_pairs:-1]
+        assert len(outputs) == len(lses), (
+            "Expected the same number of outputs as logsumexp tensors, "
+            + f"got {len(outputs)=}, {len(lses)=}"
+        )
+        merged_output, merged_lse = _merge_attentions_op(
+            outputs,  # type: ignore[arg-type]
+            lses,  # type: ignore[arg-type]
+            torch_compile=torch_compile,
+        )
+        ctx.num_pairs = num_pairs
+        ctx.save_for_backward(merged_output, merged_lse, *outputs, *lses)
+        return merged_output, merged_lse
+    @staticmethod
+    @amp_bwd
+    def backward(ctx, grad_out: Tensor, grad_lse: Tensor) -> Tuple:
+        num_pairs = ctx.num_pairs
+        merged_output, merged_lse = ctx.saved_tensors[:2]
+        outputs = ctx.saved_tensors[2 : num_pairs + 2]
+        lses = ctx.saved_tensors[num_pairs + 2 :]
+        # Outputs and LSEs from the originating attention ops must be replaced with
+        # the merged ones inplace so that we get correct behavior, and not break torch.compile
+        # graphs in the process.
+        for output, lse in zip(outputs, lses):
+            output.data.copy_(merged_output.data.reshape(output.shape))
+            lse.data.copy_(merged_lse.data.reshape(lse.shape))
+        return (
+            *(grad_out for _ in range(num_pairs)),
+            *(grad_lse for _ in range(num_pairs)),
+            None,
+        )
+def merge_attentions(
+    outputs: List[Tensor],
+    lse_tensors: List[Tensor],
+    torch_compile: bool = True,
+    use_autograd_fix: bool = True,
+) -> Tuple[Tensor, Tensor]:
+    """Takes multiple attention *outputs* originating from the same query tensor, and their
+    corresponding logsumexps, and merges them as if their context (key/value pair) had been
+    concatenated.
+    This operation is used to implement cross-neighborhood attention, and can also be used for
+    distributed setups, such as context-parallelism.
+    This operation also attempts to use `torch.compile` to fuse the elementwise operations. This
+    can be disabled by passing `torch_compile=False`.
+    Parameters:
+        outputs (List[Tensor]): List of 4-D attention output tensors, with the heads last layout
+            (`[batch, seqlen, heads, head_dim]`)
+        lse_tensors (List[Tensor]): List of 3-D logsumexp tensors, with the heads last layout
+            (`[batch, seqlen, heads]`)
+        torch_compile (bool): Attempt to use `torch.compile` to fuse the underlying elementwise
+            operations. Default: True.
+        use_autograd_fix (bool): fix backpropagation by using a custom autograd function. Only
+            compatible with fused attention operations (Flash/FMHA/FNA), only as long as the inputs
+            of this function are (views) of outputs from said attention operation.
+            NATTEN's tests (tests/test_attn_merge.py) only verify correctness for when using
+            attention operations from NATTEN. Integration for non-NATTEN ops must be verified by the
+            end user.
+            This must be disabled when using unfused Attention, which includes Flex without
+            torch.compile. Default: True.
+    Returns:
+        output (Tensor): merged attention output.
+        logsumexp (Tensor): updated logsumexp.
+    """
+    if len(outputs) < 2:
+        raise ValueError("`merge_attentions` expects at least two tensors.")
+    if len(outputs) != len(lse_tensors):
+        raise ValueError(
+            "`merge_attentions` expected number of outputs and LSE tensors to match, "
+            f"got {len(outputs)=} != {len(lse_tensors)}."
+        )
+    assert len(outputs) == len(lse_tensors)
+    requires_grad = outputs[0].requires_grad
+    shape = outputs[0].shape
+    for i, (output, lse) in enumerate(zip(outputs, lse_tensors)):
+        if output.dim() != 4 or not output.is_contiguous():
+            raise ValueError(
+                "Output tensors must be rank-4 tensors with (batch, seq, heads, dim), "
+                f"but got output {i} with rank={output.dim()}."
+            )
+        if output.shape != shape:
+            raise ValueError(
+                f"Output tensors must must match in shape, but got output {i} "
+                f"with shape={output.shape}."
+            )
+        if lse.dim() != 3:
+            raise ValueError(
+                "LSE tensors must be rank-3 tensors with (batch, seq, heads)"
+                f"but got LSE {i} with rank={lse.dim()}."
+            )
+        if lse.shape != shape[:3]:
+            raise ValueError(
+                f"LSE tensors must must match outputs in shape except last dim "
+                f"({shape=}), but got LSE {i} with shape={lse.shape}."
+            )
+        if output.requires_grad and not requires_grad:
+            raise ValueError(
+                "Either all attentions must require grad, or none of them."
+            )
+    # This path is the correct way to do backward pass, but since we can't have lists as inputs to
+    # autograd functions, we're forced to specialize it for 2-way for now.
+    if use_autograd_fix:
+        merged_output, merged_lse = MergeAttentionsAutogradFn.apply(
+            *outputs, *lse_tensors, torch_compile
+        )
+        return merged_output, merged_lse
+    return _merge_attentions_op(
+        [output.contiguous() for output in outputs],
+        [lse.contiguous() for lse in lse_tensors],
+        torch_compile=torch_compile,
+    )
+__all__ = ["merge_attentions"]

build/torch210-cu126-aarch64-linux/backends/__init__.py ADDED Viewed

	@@ -0,0 +1,254 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+from typing import List
+from ..utils import log
+logger = log.get_logger(__name__)
+import torch  # noqa: F401
+from torch import Tensor
+from ..backends.blackwell_fmha import cutlass_blackwell_fmha
+from ..backends.blackwell_fna import (
+    cutlass_blackwell_fna_generic,
+    na1d_cutlass_blackwell_fna,
+    na2d_cutlass_blackwell_fna,
+    na3d_cutlass_blackwell_fna,
+)
+from ..backends.configs import (
+    get_bwd_configs_for_cutlass_blackwell_fmha,
+    get_bwd_configs_for_cutlass_blackwell_fna,
+    get_bwd_configs_for_cutlass_fmha,
+    get_bwd_configs_for_cutlass_fna,
+    get_bwd_configs_for_cutlass_hopper_fmha,
+    get_bwd_configs_for_cutlass_hopper_fna,
+    get_configs_for_cutlass_blackwell_fmha,
+    get_configs_for_cutlass_blackwell_fna,
+    get_configs_for_cutlass_fmha,
+    get_configs_for_cutlass_fna,
+    get_configs_for_cutlass_hopper_fmha,
+    get_configs_for_cutlass_hopper_fna,
+    get_configs_for_flex_fmha,
+    get_configs_for_flex_fna,
+)
+from ..backends.configs.checks import (
+    can_run_cutlass_blackwell_fmha,
+    can_run_cutlass_blackwell_fna,
+    can_run_cutlass_fna,
+    can_run_cutlass_hopper_fmha,
+    can_run_cutlass_hopper_fna,
+    can_run_flex_attention,
+)
+from ..backends.flex import (
+    flex_fmha,
+    flex_fna_generic,
+    na1d_flex,
+    na2d_flex,
+    na3d_flex,
+)
+from ..backends.fmha import can_run_cutlass_fmha, cutlass_fmha
+from ..backends.fna import (
+    cutlass_fna_generic,
+    na1d_cutlass_fna,
+    na2d_cutlass_fna,
+    na3d_cutlass_fna,
+)
+from ..backends.hopper_fmha import cutlass_hopper_fmha
+from ..backends.hopper_fna import (
+    cutlass_hopper_fna_generic,
+    na1d_cutlass_hopper_fna,
+    na2d_cutlass_hopper_fna,
+    na3d_cutlass_hopper_fna,
+)
+def choose_backend(
+    query: Tensor, key: Tensor, value: Tensor, torch_compile: bool
+) -> str:
+    if can_run_cutlass_blackwell_fna(query, key, value):
+        logger.debug("Backend not set; picked Blackwell FNA kernel.")
+        return "blackwell-fna"
+    if can_run_cutlass_hopper_fna(query, key, value):
+        logger.debug("Backend not set; picked Hopper FNA kernel.")
+        return "hopper-fna"
+    if can_run_cutlass_fna(query, key, value):
+        logger.debug("Backend not set; picked CUTLASS (2.X) FNA kernel.")
+        return "cutlass-fna"
+    if can_run_flex_attention(query, key, value, torch_compile=torch_compile):
+        logger.debug("Backend not set; picked Flex Attention kernel.")
+        return "flex-fna"
+    raise NotImplementedError(
+        "NATTEN could not find a suitable backend for this use case. "
+        "Run with NATTEN_LOG_LEVEL=DEBUG to find out why."
+    )
+def choose_fmha_backend(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    is_causal: bool,
+    is_varlen: bool,
+    torch_compile: bool,
+) -> str:
+    if can_run_cutlass_blackwell_fmha(
+        query, key, value, is_causal=is_causal, is_varlen=is_varlen
+    ):
+        logger.debug("Backend not set; picked Blackwell FMHA kernel.")
+        return "blackwell-fmha"
+    if can_run_cutlass_hopper_fmha(
+        query, key, value, is_causal=is_causal, is_varlen=is_varlen
+    ):
+        logger.debug("Backend not set; picked Hopper FMHA kernel.")
+        return "hopper-fmha"
+    if can_run_cutlass_fmha(
+        query, key, value, is_causal=is_causal, is_varlen=is_varlen
+    ):
+        logger.debug("Backend not set; picked CUTLASS (2.X) FMHA kernel.")
+        return "cutlass-fmha"
+    if can_run_flex_attention(
+        query,
+        key,
+        value,
+        is_causal=is_causal,
+        is_varlen=is_varlen,
+        torch_compile=torch_compile,
+    ):
+        logger.debug("Backend not set; picked Flex Attention kernel.")
+        return "flex-fmha"
+    raise NotImplementedError(
+        "NATTEN could not find a suitable backend for this FMHA use case. "
+        "Run with NATTEN_LOG_LEVEL=DEBUG to find out why."
+    )
+def get_compatible_backends(
+    query: Tensor, key: Tensor, value: Tensor, torch_compile: bool
+) -> List[str]:
+    compatible_backends = []
+    if can_run_cutlass_blackwell_fna(query, key, value):
+        compatible_backends.append("blackwell-fna")
+    if can_run_cutlass_hopper_fna(query, key, value):
+        compatible_backends.append("hopper-fna")
+    if can_run_cutlass_fna(query, key, value):
+        compatible_backends.append("cutlass-fna")
+    if can_run_flex_attention(query, key, value, torch_compile=torch_compile):
+        compatible_backends.append("flex-fna")
+    return compatible_backends
+def get_compatible_fmha_backends(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    is_causal: bool,
+    is_varlen: bool,
+    torch_compile: bool,
+) -> List[str]:
+    compatible_backends = []
+    if can_run_cutlass_blackwell_fmha(
+        query, key, value, is_causal=is_causal, is_varlen=is_varlen
+    ):
+        compatible_backends.append("blackwell-fmha")
+    if can_run_cutlass_hopper_fmha(
+        query, key, value, is_causal=is_causal, is_varlen=is_varlen
+    ):
+        compatible_backends.append("hopper-fmha")
+    if can_run_cutlass_fmha(
+        query, key, value, is_causal=is_causal, is_varlen=is_varlen
+    ):
+        compatible_backends.append("cutlass-fmha")
+    if can_run_flex_attention(
+        query,
+        key,
+        value,
+        is_causal=is_causal,
+        is_varlen=is_varlen,
+        torch_compile=torch_compile,
+    ):
+        compatible_backends.append("flex-fmha")
+    return compatible_backends
+__all__ = [
+    "can_run_cutlass_fmha",
+    "can_run_cutlass_fna",
+    "can_run_cutlass_blackwell_fmha",
+    "can_run_cutlass_blackwell_fna",
+    "can_run_cutlass_hopper_fmha",
+    "can_run_cutlass_hopper_fna",
+    "can_run_flex_attention",
+    "cutlass_fmha",
+    "cutlass_fna_generic",
+    "na1d_cutlass_fna",
+    "na2d_cutlass_fna",
+    "na3d_cutlass_fna",
+    "cutlass_blackwell_fmha",
+    "cutlass_blackwell_fna_generic",
+    "cutlass_hopper_fmha",
+    "cutlass_hopper_fna_generic",
+    "na1d_cutlass_blackwell_fna",
+    "na2d_cutlass_blackwell_fna",
+    "na3d_cutlass_blackwell_fna",
+    "flex_fmha",
+    "flex_fna_generic",
+    "na1d_flex",
+    "na2d_flex",
+    "na3d_flex",
+    "na1d_cutlass_hopper_fna",
+    "na2d_cutlass_hopper_fna",
+    "na3d_cutlass_hopper_fna",
+    "get_bwd_configs_for_cutlass_fmha",
+    "get_bwd_configs_for_cutlass_fna",
+    "get_bwd_configs_for_cutlass_blackwell_fmha",
+    "get_bwd_configs_for_cutlass_blackwell_fna",
+    "get_configs_for_cutlass_blackwell_fmha",
+    "get_configs_for_cutlass_blackwell_fna",
+    "get_configs_for_cutlass_fmha",
+    "get_configs_for_cutlass_fna",
+    "get_configs_for_cutlass_hopper_fmha",
+    "get_bwd_configs_for_cutlass_hopper_fmha",
+    "get_configs_for_cutlass_hopper_fna",
+    "get_bwd_configs_for_cutlass_hopper_fna",
+    "get_configs_for_flex_fmha",
+    "get_configs_for_flex_fna",
+]

build/torch210-cu126-aarch64-linux/backends/blackwell_fmha.py ADDED Viewed

	@@ -0,0 +1,254 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+import functools
+from typing import Optional, Tuple, Union
+import torch
+from torch import Tensor
+from torch.amp import custom_bwd, custom_fwd
+from torch.autograd import Function
+amp_fwd = functools.partial(custom_fwd, device_type="cuda")
+amp_bwd = functools.partial(custom_bwd, device_type="cuda")
+from .._libnatten import blackwell_fmha_backward, blackwell_fmha_forward
+from ..backends.configs.checks import can_run_cutlass_blackwell_fmha
+from ..backends.configs.cutlass_blackwell import (
+    check_cutlass_blackwell_fmha_backward_config,
+    check_cutlass_blackwell_fmha_forward_config,
+)
+from .._types import (
+    CutlassBlackwellFmhaBackwardConfigType,
+    CutlassBlackwellFmhaForwardConfigType,
+    NoneType,
+)
+from ..utils import log
+from ..utils.checks import fmha_tensor_checks, varlen_tensor_checks
+logger = log.get_logger(__name__)
+class CutlassBlackwellFmhaAutogradFn(Function):
+    @staticmethod
+    @amp_fwd
+    def forward(
+        ctx,
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        is_causal: bool,
+        scale: float,
+        forward_config: CutlassBlackwellFmhaForwardConfigType,
+        backward_config: CutlassBlackwellFmhaBackwardConfigType,
+        run_persistent_kernel: bool,
+        cumulative_seqlen_Q: Optional[Tensor],
+        cumulative_seqlen_KV: Optional[Tensor],
+        max_seqlen_Q: int,
+        max_seqlen_KV: int,
+    ) -> Tuple[Tensor, Tensor]:
+        query = query.contiguous()
+        key = key.contiguous()
+        value = value.contiguous()
+        q_tile_size, kv_tile_size = forward_config
+        output, logsumexp = blackwell_fmha_forward(
+            query,
+            key,
+            value,
+            is_causal,
+            scale,
+            q_tile_size,
+            kv_tile_size,
+            run_persistent_kernel,
+            cumulative_seqlen_Q,
+            cumulative_seqlen_KV,
+            max_seqlen_Q,
+            max_seqlen_KV,
+        )
+        ctx.save_for_backward(
+            query,
+            key,
+            value,
+            logsumexp,
+            output,
+            cumulative_seqlen_Q,
+            cumulative_seqlen_KV,
+        )
+        ctx.scale = scale
+        ctx.is_causal = is_causal
+        ctx.max_seqlen_Q = max_seqlen_Q
+        ctx.max_seqlen_KV = max_seqlen_KV
+        ctx.backward_config = backward_config
+        # Always record determinism behavior during forward pass (forward pass itself is
+        # deterministic anyway).
+        # Determinism could be limited to part of the program, which means during forward pass
+        # it'll be true, but on .backward() call, if it's been turned off, it will stay off when we
+        # get to this operation's backward call.
+        ctx.deterministic = torch.are_deterministic_algorithms_enabled()
+        return output, logsumexp
+    @staticmethod
+    @amp_bwd
+    def backward(ctx, grad_out: Tensor, grad_lse: Tensor) -> Tuple[
+        Tensor,
+        Tensor,
+        Tensor,
+        NoneType,
+        NoneType,
+        NoneType,
+        NoneType,
+        NoneType,
+        # varlen
+        NoneType,
+        NoneType,
+        NoneType,
+        NoneType,
+    ]:
+        (
+            query,
+            key,
+            value,
+            logsumexp,
+            output,
+            cumulative_seqlen_Q,
+            cumulative_seqlen_KV,
+        ) = ctx.saved_tensors
+        d_output = grad_out.contiguous()  # noqa: F841
+        q_tile_size, k_tile_size = ctx.backward_config
+        d_query, d_key, d_value = blackwell_fmha_backward(
+            query,
+            key,
+            value,
+            output,
+            d_output,
+            logsumexp,
+            ctx.is_causal,
+            ctx.scale,
+            q_tile_size,
+            k_tile_size,
+            cumulative_seqlen_Q,
+            cumulative_seqlen_KV,
+            ctx.max_seqlen_Q,
+            ctx.max_seqlen_KV,
+            ctx.deterministic,
+        )
+        return (
+            d_query,
+            d_key,
+            d_value,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+def cutlass_blackwell_fmha(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    is_causal: bool = False,
+    scale: Optional[float] = None,
+    q_tile_size: Optional[int] = None,
+    kv_tile_size: Optional[int] = None,
+    run_persistent_kernel: bool = False,
+    backward_q_tile_size: Optional[int] = None,
+    backward_kv_tile_size: Optional[int] = None,
+    return_lse: bool = False,
+    # varlen parameters
+    cumulative_seqlen_Q: Optional[Tensor] = None,
+    cumulative_seqlen_KV: Optional[Tensor] = None,
+    max_seqlen_Q: int = 0,
+    max_seqlen_KV: int = 0,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    fmha_tensor_checks(
+        query,
+        key,
+        value,
+        must_match_head_dims=True,
+        supports_gqa_mqa=True,
+        backend_name="Blackwell FMHA",
+    )
+    (
+        cumulative_seqlen_Q,
+        cumulative_seqlen_KV,
+        max_seqlen_Q,
+        max_seqlen_KV,
+    ) = varlen_tensor_checks(
+        query=query,
+        key=key,
+        value=value,
+        cumulative_seqlen_Q=cumulative_seqlen_Q,
+        cumulative_seqlen_KV=cumulative_seqlen_KV,
+        max_seqlen_Q=max_seqlen_Q,
+        max_seqlen_KV=max_seqlen_KV,
+    )
+    is_varlen = cumulative_seqlen_Q is not None
+    assert can_run_cutlass_blackwell_fmha(
+        query, key, value, is_causal=is_causal, is_varlen=is_varlen, raise_error=True
+    )
+    forward_config = check_cutlass_blackwell_fmha_forward_config(
+        input_tensor=query, q_tile_size=q_tile_size, kv_tile_size=kv_tile_size
+    )
+    backward_config = check_cutlass_blackwell_fmha_backward_config(
+        input_tensor=query,
+        q_tile_size=backward_q_tile_size,
+        kv_tile_size=backward_kv_tile_size,
+    )
+    scale = scale or query.shape[-1] ** -0.5
+    output, lse = CutlassBlackwellFmhaAutogradFn.apply(
+        query,
+        key,
+        value,
+        is_causal,
+        scale,
+        forward_config,
+        backward_config,
+        run_persistent_kernel,
+        cumulative_seqlen_Q,
+        cumulative_seqlen_KV,
+        max_seqlen_Q,
+        max_seqlen_KV,
+    )
+    if return_lse:
+        return output, lse
+    return output

build/torch210-cu126-aarch64-linux/backends/blackwell_fna.py ADDED Viewed

	@@ -0,0 +1,500 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+import functools
+from typing import Optional, Tuple, Union
+import torch
+from torch import Tensor
+from torch.amp import custom_bwd, custom_fwd
+from torch.autograd import Function
+amp_fwd = functools.partial(custom_fwd, device_type="cuda")
+amp_bwd = functools.partial(custom_bwd, device_type="cuda")
+from .._libnatten import (
+    blackwell_na1d_backward,
+    blackwell_na1d_forward,
+    blackwell_na2d_backward,
+    blackwell_na2d_forward,
+    blackwell_na3d_backward,
+    blackwell_na3d_forward,
+)
+from ..backends.configs.checks import can_run_cutlass_blackwell_fna
+from ..backends.configs.cutlass_blackwell import (
+    check_cutlass_blackwell_fna_backward_config,
+    check_cutlass_blackwell_fna_forward_config,
+)
+from ..token_permute import token_permute_operation, token_unpermute_operation
+from .._types import (
+    CausalArg1DTypeOrDed,
+    CausalArg2DTypeOrDed,
+    CausalArg3DTypeOrDed,
+    CausalArgType,
+    CausalArgTypeOrDed,
+    CutlassBlackwellFnaBackwardConfigType,
+    CutlassBlackwellFnaForwardConfigType,
+    Dimension1DType,
+    Dimension1DTypeOrDed,
+    Dimension2DType,
+    Dimension2DTypeOrDed,
+    Dimension3DType,
+    Dimension3DTypeOrDed,
+    DimensionType,
+    DimensionTypeOrDed,
+    NoneType,
+)
+from ..utils.checks import (
+    check_all_args,
+    check_args_against_input,
+    na_tensor_checks,
+)
+def make_cutlass_blackwell_fna_autograd_fn(na_dim):
+    assert na_dim in [1, 2, 3]
+    FORWARD_OPS = {
+        1: blackwell_na1d_forward,
+        2: blackwell_na2d_forward,
+        3: blackwell_na3d_forward,
+    }
+    BACKWARD_OPS = {
+        1: blackwell_na1d_backward,
+        2: blackwell_na2d_backward,
+        3: blackwell_na3d_backward,
+    }
+    class CutlassBlackwellFnaGenericAutogradFn(Function):
+        @staticmethod
+        @amp_fwd
+        def forward(
+            ctx,
+            query: Tensor,
+            key: Tensor,
+            value: Tensor,
+            kernel_size: DimensionType,
+            stride: DimensionType,
+            dilation: DimensionType,
+            is_causal: CausalArgType,
+            scale: float,
+            forward_config: CutlassBlackwellFnaForwardConfigType,
+            backward_config: CutlassBlackwellFnaBackwardConfigType,
+            run_persistent_kernel: bool,
+        ) -> Tuple[Tensor, Tensor]:
+            kernel_size, stride, dilation, is_causal = check_all_args(
+                na_dim, kernel_size, stride, dilation, is_causal
+            )
+            q_tile_shape, kv_tile_shape = forward_config
+            # Token permute begin
+            query_perm, qkv_shape, q_shape = token_permute_operation(
+                query,
+                tile_shape=q_tile_shape,
+                dilation=dilation,
+                flip_tiled_dims=True,
+            )
+            key_perm, _, k_shape = token_permute_operation(
+                key, tile_shape=kv_tile_shape, dilation=dilation, flip_tiled_dims=True
+            )
+            value_perm, _, v_shape = token_permute_operation(
+                value, tile_shape=kv_tile_shape, dilation=dilation, flip_tiled_dims=True
+            )
+            assert k_shape == v_shape
+            kv_shape = k_shape
+            # Token permute end
+            query_perm = query_perm.contiguous()
+            key_perm = key_perm.contiguous()
+            value_perm = value_perm.contiguous()
+            output_perm, logsumexp_perm = FORWARD_OPS[na_dim](
+                query_perm,
+                key_perm,
+                value_perm,
+                kernel_size,
+                stride,
+                dilation,
+                is_causal,
+                scale,
+                q_shape,
+                kv_shape,
+                qkv_shape,
+                q_tile_shape,
+                kv_tile_shape,
+                run_persistent_kernel,
+            )
+            # Token un-permute begin
+            output = token_unpermute_operation(
+                output_perm,
+                token_layout_shape=qkv_shape,
+                tile_shape=q_tile_shape,
+                dilation=dilation,
+                flip_tiled_dims=True,
+            )
+            logsumexp = token_unpermute_operation(
+                logsumexp_perm.unsqueeze(-1),
+                token_layout_shape=qkv_shape,
+                tile_shape=q_tile_shape,
+                dilation=dilation,
+                flip_tiled_dims=True,
+            ).squeeze(-1)
+            # Token un-permute end
+            ctx.save_for_backward(query, key, value, logsumexp, output)
+            ctx.kernel_size = kernel_size
+            ctx.stride = stride
+            ctx.dilation = dilation
+            ctx.is_causal = is_causal
+            ctx.scale = scale
+            ctx.backward_config = backward_config
+            # Always record determinism behavior during forward pass (forward pass itself is
+            # deterministic anyway).
+            # Determinism could be limited to part of the program, which means during forward pass
+            # it'll be true, but on .backward() call, if it's been turned off, it will stay off when we
+            # get to this operation's backward call.
+            ctx.deterministic = torch.are_deterministic_algorithms_enabled()
+            return output, logsumexp
+        @staticmethod
+        @amp_bwd
+        def backward(ctx, d_output: Tensor, d_lse: Tensor) -> Tuple[
+            Tensor,
+            Tensor,
+            Tensor,
+            NoneType,
+            NoneType,
+            NoneType,
+            NoneType,
+            NoneType,
+            NoneType,
+            NoneType,
+            NoneType,
+        ]:
+            query, key, value, logsumexp, output = ctx.saved_tensors
+            kernel_size, stride, dilation, is_causal, scale = (
+                ctx.kernel_size,
+                ctx.stride,
+                ctx.dilation,
+                ctx.is_causal,
+                ctx.scale,
+            )
+            q_tile_shape, kv_tile_shape = ctx.backward_config
+            if ctx.deterministic:
+                raise RuntimeError(
+                    "Blackwell FNA backward pass does not have a deterministic mode, "
+                    "but PyTorch's deterministic algorithms were enabled. To proceed, "
+                    "you must either disable torch's deterministic mode, or choose a "
+                    "different backend."
+                )
+            # Token permute begin
+            query_perm, qkv_shape, q_shape = token_permute_operation(
+                query, tile_shape=q_tile_shape, dilation=dilation, flip_tiled_dims=True
+            )
+            output_perm, _, o_shape = token_permute_operation(
+                output, tile_shape=q_tile_shape, dilation=dilation, flip_tiled_dims=True
+            )
+            d_output_perm, _, d_o_shape = token_permute_operation(
+                d_output,
+                tile_shape=q_tile_shape,
+                dilation=dilation,
+                flip_tiled_dims=True,
+            )
+            logsumexp_perm, _, _ = token_permute_operation(
+                logsumexp.unsqueeze(-1),
+                tile_shape=q_tile_shape,
+                dilation=dilation,
+                flip_tiled_dims=True,
+            )
+            key_perm, _, k_shape = token_permute_operation(
+                key, tile_shape=kv_tile_shape, dilation=dilation, flip_tiled_dims=True
+            )
+            value_perm, _, v_shape = token_permute_operation(
+                value, tile_shape=kv_tile_shape, dilation=dilation, flip_tiled_dims=True
+            )
+            assert q_shape == o_shape == d_o_shape
+            assert k_shape == v_shape
+            kv_shape = k_shape
+            # Token permute end
+            query_perm = query_perm.contiguous()
+            key_perm = key_perm.contiguous()
+            value_perm = value_perm.contiguous()
+            output_perm = output_perm.contiguous()
+            d_output_perm = d_output_perm.contiguous()
+            logsumexp_perm = logsumexp_perm.squeeze(-1)
+            d_query_perm, d_key_perm, d_value_perm = BACKWARD_OPS[na_dim](
+                query_perm,
+                key_perm,
+                value_perm,
+                output_perm,
+                d_output_perm,
+                logsumexp_perm,
+                kernel_size,
+                stride,
+                dilation,
+                is_causal,
+                scale,
+                q_shape,
+                kv_shape,
+                qkv_shape,
+                q_tile_shape,
+                kv_tile_shape,
+            )
+            # Token un-permute begin
+            d_query = token_unpermute_operation(
+                d_query_perm,
+                token_layout_shape=qkv_shape,
+                tile_shape=q_tile_shape,
+                dilation=dilation,
+                flip_tiled_dims=True,
+            )
+            d_key = token_unpermute_operation(
+                d_key_perm,
+                token_layout_shape=qkv_shape,
+                tile_shape=kv_tile_shape,
+                dilation=dilation,
+                flip_tiled_dims=True,
+            )
+            d_value = token_unpermute_operation(
+                d_value_perm,
+                token_layout_shape=qkv_shape,
+                tile_shape=kv_tile_shape,
+                dilation=dilation,
+                flip_tiled_dims=True,
+            )
+            # Token un-permute end
+            assert d_query.shape == query.shape
+            assert d_key.shape == key.shape
+            assert d_value.shape == value.shape
+            return (
+                d_query,
+                d_key,
+                d_value,
+                None,
+                None,
+                None,
+                None,
+                None,
+                None,
+                None,
+                None,
+            )
+    return CutlassBlackwellFnaGenericAutogradFn
+CutlassBlackwellFna1DAutogradFn = make_cutlass_blackwell_fna_autograd_fn(1)
+CutlassBlackwellFna2DAutogradFn = make_cutlass_blackwell_fna_autograd_fn(2)
+CutlassBlackwellFna3DAutogradFn = make_cutlass_blackwell_fna_autograd_fn(3)
+CutlassBlackwellFNAAutogradFns = {
+    1: CutlassBlackwellFna1DAutogradFn,
+    2: CutlassBlackwellFna2DAutogradFn,
+    3: CutlassBlackwellFna3DAutogradFn,
+}
+def cutlass_blackwell_fna_generic(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    kernel_size: DimensionTypeOrDed,
+    stride: DimensionTypeOrDed = 1,
+    dilation: DimensionTypeOrDed = 1,
+    is_causal: Optional[CausalArgTypeOrDed] = False,
+    scale: Optional[float] = None,
+    q_tile_shape: Optional[DimensionType] = None,
+    kv_tile_shape: Optional[DimensionType] = None,
+    backward_q_tile_shape: Optional[DimensionType] = None,
+    backward_kv_tile_shape: Optional[DimensionType] = None,
+    run_persistent_kernel: bool = False,
+    return_lse: bool = False,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    na_tensor_checks(
+        query, key, value, must_match_head_dims=True, supports_gqa_mqa=True
+    )
+    assert can_run_cutlass_blackwell_fna(query, key, value, raise_error=True)
+    na_dim = query.dim() - 3  # batch, heads, head_dim
+    kernel_size, stride, dilation, is_causal = check_all_args(
+        na_dim, kernel_size, stride, dilation, is_causal
+    )
+    check_args_against_input(
+        query,
+        kernel_size=kernel_size,
+        stride=stride,
+        dilation=dilation,
+        is_causal=is_causal,
+    )
+    forward_config = check_cutlass_blackwell_fna_forward_config(
+        input_tensor=query, q_tile_shape=q_tile_shape, kv_tile_shape=kv_tile_shape
+    )
+    backward_config = check_cutlass_blackwell_fna_backward_config(
+        input_tensor=query,
+        q_tile_shape=backward_q_tile_shape,
+        kv_tile_shape=backward_kv_tile_shape,
+    )
+    scale = scale or query.shape[-1] ** -0.5
+    output, lse = CutlassBlackwellFNAAutogradFns[na_dim].apply(
+        query,
+        key,
+        value,
+        kernel_size,
+        stride,
+        dilation,
+        is_causal,
+        scale,
+        forward_config,
+        backward_config,
+        run_persistent_kernel,
+    )
+    if return_lse:
+        return output, lse
+    return output
+def na1d_cutlass_blackwell_fna(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    kernel_size: Dimension1DTypeOrDed,
+    stride: Dimension1DTypeOrDed = 1,
+    dilation: Dimension1DTypeOrDed = 1,
+    is_causal: Optional[CausalArg1DTypeOrDed] = False,
+    scale: Optional[float] = None,
+    q_tile_shape: Optional[Dimension1DType] = None,
+    kv_tile_shape: Optional[Dimension1DType] = None,
+    backward_q_tile_shape: Optional[Dimension1DType] = None,
+    backward_kv_tile_shape: Optional[Dimension1DType] = None,
+    run_persistent_kernel: bool = False,
+    return_lse: bool = False,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    return cutlass_blackwell_fna_generic(
+        query=query,
+        key=key,
+        value=value,
+        kernel_size=kernel_size,
+        stride=stride,
+        dilation=dilation,
+        is_causal=is_causal,
+        scale=scale,
+        q_tile_shape=q_tile_shape,
+        kv_tile_shape=kv_tile_shape,
+        backward_q_tile_shape=backward_q_tile_shape,
+        backward_kv_tile_shape=backward_kv_tile_shape,
+        run_persistent_kernel=run_persistent_kernel,
+        return_lse=return_lse,
+    )
+def na2d_cutlass_blackwell_fna(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    kernel_size: Dimension2DTypeOrDed,
+    stride: Dimension2DTypeOrDed = 1,
+    dilation: Dimension2DTypeOrDed = 1,
+    is_causal: Optional[CausalArg2DTypeOrDed] = False,
+    scale: Optional[float] = None,
+    q_tile_shape: Optional[Dimension2DType] = None,
+    kv_tile_shape: Optional[Dimension2DType] = None,
+    backward_q_tile_shape: Optional[Dimension2DType] = None,
+    backward_kv_tile_shape: Optional[Dimension2DType] = None,
+    run_persistent_kernel: bool = False,
+    return_lse: bool = False,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    return cutlass_blackwell_fna_generic(
+        query=query,
+        key=key,
+        value=value,
+        kernel_size=kernel_size,
+        stride=stride,
+        dilation=dilation,
+        is_causal=is_causal,
+        scale=scale,
+        q_tile_shape=q_tile_shape,
+        kv_tile_shape=kv_tile_shape,
+        backward_q_tile_shape=backward_q_tile_shape,
+        backward_kv_tile_shape=backward_kv_tile_shape,
+        run_persistent_kernel=run_persistent_kernel,
+        return_lse=return_lse,
+    )
+def na3d_cutlass_blackwell_fna(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    kernel_size: Dimension3DTypeOrDed,
+    stride: Dimension3DTypeOrDed = 1,
+    dilation: Dimension3DTypeOrDed = 1,
+    is_causal: Optional[CausalArg3DTypeOrDed] = False,
+    scale: Optional[float] = None,
+    q_tile_shape: Optional[Dimension3DType] = None,
+    kv_tile_shape: Optional[Dimension3DType] = None,
+    backward_q_tile_shape: Optional[Dimension3DType] = None,
+    backward_kv_tile_shape: Optional[Dimension3DType] = None,
+    run_persistent_kernel: bool = False,
+    return_lse: bool = False,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    return cutlass_blackwell_fna_generic(
+        query=query,
+        key=key,
+        value=value,
+        kernel_size=kernel_size,
+        stride=stride,
+        dilation=dilation,
+        is_causal=is_causal,
+        scale=scale,
+        q_tile_shape=q_tile_shape,
+        kv_tile_shape=kv_tile_shape,
+        backward_q_tile_shape=backward_q_tile_shape,
+        backward_kv_tile_shape=backward_kv_tile_shape,
+        run_persistent_kernel=run_persistent_kernel,
+        return_lse=return_lse,
+    )

build/torch210-cu126-aarch64-linux/backends/configs/__init__.py ADDED Viewed

	@@ -0,0 +1,584 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+from typing import List
+from ...utils import log
+logger = log.get_logger(__name__)
+import torch  # noqa: F401
+from torch import Tensor
+from ...backends.configs.checks import (
+    can_run_cutlass_blackwell_fmha,
+    can_run_cutlass_blackwell_fna,
+    can_run_cutlass_fmha,
+    can_run_cutlass_fna,
+    can_run_cutlass_hopper_fmha,
+    can_run_cutlass_hopper_fna,
+    can_run_flex_attention,
+)
+from ...backends.configs.cutlass import (
+    get_all_tile_shapes_backward as get_all_cutlass_fna_backward_configs,
+    get_all_tile_shapes_forward as get_all_cutlass_fna_forward_configs,
+    get_all_tile_sizes_backward as get_all_cutlass_fmha_backward_configs,
+    get_all_tile_sizes_forward as get_all_cutlass_fmha_forward_configs,
+)
+from ...backends.configs.cutlass_blackwell import (
+    get_all_backward_configs as get_all_blackwell_fna_backward_configs,
+    get_all_fmha_backward_configs as get_all_blackwell_fmha_backward_configs,
+    get_all_fmha_forward_configs as get_all_blackwell_fmha_forward_configs,
+    get_all_forward_configs as get_all_blackwell_fna_forward_configs,
+)
+from ...backends.configs.cutlass_hopper import (
+    get_all_backward_configs as get_all_hopper_fna_backward_configs,
+    get_all_fmha_backward_configs as get_all_hopper_fmha_backward_configs,
+    get_all_fmha_forward_configs as get_all_hopper_fmha_forward_configs,
+    get_all_forward_configs as get_all_hopper_fna_forward_configs,
+)
+from ...backends.configs.flex import (
+    get_all_tile_shapes_forward as get_all_flex_fna_forward_configs,
+    get_all_tile_sizes_forward as get_all_flex_fmha_forward_configs,
+)
+from ..._types import (
+    CutlassBlackwellFmhaBackwardConfigType,
+    CutlassBlackwellFmhaForwardConfigType,
+    CutlassBlackwellFnaBackwardConfigType,
+    CutlassBlackwellFnaForwardConfigType,
+    CutlassFmhaBackwardConfigType,
+    CutlassFmhaForwardConfigType,
+    CutlassFnaBackwardConfigType,
+    CutlassFnaForwardConfigType,
+    CutlassHopperFmhaBackwardConfigType,
+    CutlassHopperFmhaForwardConfigType,
+    CutlassHopperFnaBackwardConfigType,
+    CutlassHopperFnaForwardConfigType,
+    FlexFmhaForwardConfigType,
+    FlexFnaForwardConfigType,
+)
+### CUTLASS Blackwell kernels
+def get_configs_for_cutlass_blackwell_fmha(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+) -> List[CutlassBlackwellFmhaForwardConfigType]:
+    """Returns Blackwell FMHA configurations compatible with input tensors, if any.
+    Checks first if a CUDA tensor, and on a Blackwell datacenter GPU (SM100; compute capability
+    10.0), and if so, returns *forward pass* configurations compatible with the tensor dtype and
+    head dim.
+    Each configuration for this operation is a tuple of two integers: `(q_tile_size,
+    kv_tile_size)`. These are arguments to [natten.attention][natten.attention].
+    Args:
+        query: Query tensor matching the shape, dtype, and device of your use case.
+        key:   Key tensor matching the shape, dtype, and device of your use case.
+        value: Value tensor matching the shape, dtype, and device of your use case.
+    Returns:
+        (List[Tuple[int, int]]): List of tuples of two integers corresponding to query and KV tile
+            sizes.
+    """
+    if not can_run_cutlass_blackwell_fmha(
+        query=query,
+        key=key,
+        value=value,
+        is_causal=False,
+        is_varlen=False,
+        raise_error=False,
+    ):
+        return []
+    return get_all_blackwell_fmha_forward_configs(input_tensor=query)
+def get_bwd_configs_for_cutlass_blackwell_fmha(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+) -> List[CutlassBlackwellFmhaBackwardConfigType]:
+    """Returns Blackwell FMHA backward pass configurations compatible with input tensors, if any.
+    Checks first if a CUDA tensor, and on a Blackwell datacenter GPU (SM100; compute capability
+    10.0), and if so, returns *backward pass* configurations compatible with the tensor dtype and
+    head dim.
+    Each configuration for this operation is a tuple of two integers: `(q_tile_size,
+    kv_tile_size)`. These are arguments to [natten.attention][natten.attention].
+    Args:
+        query: Query tensor matching the shape, dtype, and device of your use case.
+        key:   Key tensor matching the shape, dtype, and device of your use case.
+        value: Value tensor matching the shape, dtype, and device of your use case.
+    Returns:
+        (List[Tuple[int, int]]): List of tuples of two integers corresponding to query and KV tile
+            sizes.
+    """
+    if not can_run_cutlass_blackwell_fmha(
+        query=query,
+        key=key,
+        value=value,
+        is_causal=False,
+        is_varlen=False,
+        raise_error=False,
+    ):
+        return []
+    return get_all_blackwell_fmha_backward_configs(input_tensor=query)
+def get_configs_for_cutlass_blackwell_fna(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+) -> List[CutlassBlackwellFnaForwardConfigType]:
+    """Returns Blackwell FNA configurations compatible with input tensors, if any.
+    Checks first if a CUDA tensor, and on a Blackwell datacenter GPU (SM100; compute capability
+    10.0), and if so, returns *forward pass* configurations compatible with the tensor dtype and
+    head dim, and according to the rank of the token layout (1D/2D/3D).
+    Each configuration for this operation is a tuple of two integer tuples: `(q_tile_shape,
+    kv_tile_shape)`. These are arguments to [natten.na1d][natten.na1d], [natten.na2d][natten.na2d],
+    and [natten.na3d][natten.na3d].
+    Args:
+        query: Query tensor matching the shape, dtype, and device of your use case.
+        key:   Key tensor matching the shape, dtype, and device of your use case.
+        value: Value tensor matching the shape, dtype, and device of your use case.
+    Returns:
+        (List[Tuple[tuple, tuple]]): List of tuples of two integer tuples corresponding to query
+            and KV tile *shapes*.
+    """
+    if not can_run_cutlass_blackwell_fna(
+        query=query, key=key, value=value, raise_error=False
+    ):
+        return []
+    return get_all_blackwell_fna_forward_configs(input_tensor=query)
+def get_bwd_configs_for_cutlass_blackwell_fna(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+) -> List[CutlassBlackwellFnaBackwardConfigType]:
+    """Returns Blackwell FNA backward pass configurations compatible with input tensors, if any.
+    Checks first if a CUDA tensor, and on a Blackwell datacenter GPU (SM100; compute capability
+    10.0), and if so, returns *backward pass* configurations compatible with the tensor dtype and
+    head dim, and according to the rank of the token layout (1D/2D/3D).
+    Each configuration for this operation is a tuple of two integer tuples: `(q_tile_shape,
+    kv_tile_shape)`. These are arguments to [natten.na1d][natten.na1d], [natten.na2d][natten.na2d],
+    and [natten.na3d][natten.na3d].
+    Args:
+        query: Query tensor matching the shape, dtype, and device of your use case.
+        key:   Key tensor matching the shape, dtype, and device of your use case.
+        value: Value tensor matching the shape, dtype, and device of your use case.
+    Returns:
+        (List[Tuple[tuple, tuple]]): List of tuples of two integer tuples corresponding to query
+            and KV tile *shapes*.
+    """
+    if not can_run_cutlass_blackwell_fna(
+        query=query, key=key, value=value, raise_error=False
+    ):
+        return []
+    return get_all_blackwell_fna_backward_configs(input_tensor=query)
+### CUTLASS Hopper kernels
+def get_configs_for_cutlass_hopper_fmha(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+) -> List[CutlassHopperFmhaForwardConfigType]:
+    """Returns Hopper FMHA configurations compatible with input tensors, if any.
+    Checks first if a CUDA tensor, and on a Hopper GPU (SM90; compute capability 9.0), and if so,
+    returns *forward pass* configurations compatible with the tensor dtype and head dim.
+    Each configuration for this operation is a tuple of one integer tuple, and another integer:
+    `((q_tile_size, kv_tile_size), kernel_schedule)`. These are arguments to
+    [natten.attention][natten.attention].
+    `kernel_schedule` is specific to Hopper FNA/FMHA only.
+    Args:
+        query: Query tensor matching the shape, dtype, and device of your use case.
+        key:   Key tensor matching the shape, dtype, and device of your use case.
+        value: Value tensor matching the shape, dtype, and device of your use case.
+    Returns:
+        (List[Tuple[Tuple[int, int], KernelSchedule]]): List of tuples of one tuple of two integers
+            corresponding to query and KV tile sizes, and a kernel schedule enum type.
+    """
+    if not can_run_cutlass_hopper_fmha(
+        query=query,
+        key=key,
+        value=value,
+        is_causal=False,
+        is_varlen=False,
+        raise_error=False,
+    ):
+        return []
+    return get_all_hopper_fmha_forward_configs(input_tensor=query)
+def get_bwd_configs_for_cutlass_hopper_fmha(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+) -> List[CutlassHopperFmhaBackwardConfigType]:
+    """Returns Hopper FMHA backward pass configurations compatible with input tensors, if any.
+    Checks first if a CUDA tensor, and on a Hopper GPU (SM90; compute capability 9.0), and if so,
+    returns *backward pass* configurations compatible with the tensor dtype and head dim.
+    Each configuration for this operation is an integer tuple:
+    `(backward_q_tile_size, backward_kv_tile_size)`. These are arguments to
+    [natten.attention][natten.attention].
+    Note that unlike forward pass, kernel schedule is not part of the configuration. All backward
+    pass kernels are persistent warp-specialized.
+    See CUTLASS's [example 88](https://github.com/NVIDIA/cutlass/tree/main/examples/88_hopper_fmha)
+    for more.
+    Args:
+        query: Query tensor matching the shape, dtype, and device of your use case.
+        key:   Key tensor matching the shape, dtype, and device of your use case.
+        value: Value tensor matching the shape, dtype, and device of your use case.
+    Returns:
+        (List[Tuple[int, int]]): List of integer tuples corresponding to query and KV tile sizes.
+    """
+    if not can_run_cutlass_hopper_fmha(
+        query=query,
+        key=key,
+        value=value,
+        is_causal=False,
+        is_varlen=False,
+        raise_error=False,
+    ):
+        return []
+    return get_all_hopper_fmha_backward_configs(input_tensor=query)
+def get_configs_for_cutlass_hopper_fna(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+) -> List[CutlassHopperFnaForwardConfigType]:
+    """Returns Hopper FNA configurations compatible with input tensors, if any.
+    Checks first if a CUDA tensor, and on a Hopper GPU (SM90; compute capability 9.0), and if so,
+    returns *forward pass* configurations compatible with the tensor dtype and head dim.
+    Each configuration for this operation is a tuple of one tuple, and another integer:
+    `((q_tile_shape, kv_tile_shape), kernel_schedule)`. These are arguments to
+    [natten.na1d][natten.na1d], [natten.na2d][natten.na2d], and [natten.na3d][natten.na3d].
+    `kernel_schedule` is specific to Hopper FNA/FMHA only.
+    Args:
+        query: Query tensor matching the shape, dtype, and device of your use case.
+        key:   Key tensor matching the shape, dtype, and device of your use case.
+        value: Value tensor matching the shape, dtype, and device of your use case.
+    Returns:
+        (List[Tuple[Tuple[tuple, tuple], KernelSchedule]]): List of tuples of one tuple of two
+            shape tuples, corresponding to query and KV tile *shapes*, and a kernel schedule enum
+            type.
+    """
+    if not can_run_cutlass_hopper_fna(
+        query=query, key=key, value=value, raise_error=False
+    ):
+        return []
+    return get_all_hopper_fna_forward_configs(input_tensor=query)
+def get_bwd_configs_for_cutlass_hopper_fna(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+) -> List[CutlassHopperFnaBackwardConfigType]:
+    """Returns Hopper FNA backward pass configurations compatible with input tensors, if any.
+    Checks first if a CUDA tensor, and on a Hopper GPU (SM90; compute capability 9.0), and if so,
+    returns *backward pass* configurations compatible with the tensor dtype and head dim.
+    Each configuration for this operation is a tuple of two tuples:
+    `(q_tile_shape, kv_tile_shape)`. These are arguments to [natten.na1d][natten.na1d],
+    [natten.na2d][natten.na2d], and [natten.na3d][natten.na3d].
+    Note that unlike forward pass, kernel schedule is not part of the configuration. All backward
+    pass kernels are persistent warp-specialized.
+    See CUTLASS's [example 88](https://github.com/NVIDIA/cutlass/tree/main/examples/88_hopper_fmha)
+    for more.
+    Args:
+        query: Query tensor matching the shape, dtype, and device of your use case.
+        key:   Key tensor matching the shape, dtype, and device of your use case.
+        value: Value tensor matching the shape, dtype, and device of your use case.
+    Returns:
+        (List[Tuple[tuple, tuple]]): List of tuples of two shape tuples, corresponding to query and
+        KV tile *shapes*.
+    """
+    if not can_run_cutlass_hopper_fna(
+        query=query, key=key, value=value, raise_error=False
+    ):
+        return []
+    return get_all_hopper_fna_backward_configs(input_tensor=query)
+### CUTLASS 2.X kernels
+def get_configs_for_cutlass_fmha(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+) -> List[CutlassFmhaForwardConfigType]:
+    """Returns CUTLASS FMHA configurations compatible with input tensors, if any.
+    Checks first if a CUDA tensor, and on a device with compute capability >= 5.0, and if so,
+    returns *forward pass* configurations compatible with the specific compute capability, tensor
+    dtype and head dim.
+    Each configuration for this operation is a tuple of two integers: `(q_tile_size,
+    kv_tile_size)`. These are arguments to [natten.attention][natten.attention].
+    Args:
+        query: Query tensor matching the shape, dtype, and device of your use case.
+        key:   Key tensor matching the shape, dtype, and device of your use case.
+        value: Value tensor matching the shape, dtype, and device of your use case.
+    Returns:
+        (List[Tuple[int, int]]): List of tuples of two integers corresponding to query and KV tile
+            sizes.
+    """
+    if not can_run_cutlass_fmha(
+        query=query,
+        key=key,
+        value=value,
+        is_causal=False,
+        is_varlen=False,
+        raise_error=False,
+    ):
+        return []
+    return get_all_cutlass_fmha_forward_configs(input_tensor=query)
+def get_bwd_configs_for_cutlass_fmha(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+) -> List[CutlassFmhaBackwardConfigType]:
+    """Returns CUTLASS FMHA backward pass configurations compatible with input tensors, if any.
+    Checks first if a CUDA tensor, and on a device with compute capability >= 5.0, and if so,
+    returns *backward pass* configurations compatible with the specific compute capability, tensor
+    dtype and head dim.
+    Each configuration for this operation is a tuple of two integers: `(backward_q_tile_size,
+    backward_kv_tile_size)`. These are arguments to [natten.attention][natten.attention].
+    Args:
+        query: Query tensor matching the shape, dtype, and device of your use case.
+        key:   Key tensor matching the shape, dtype, and device of your use case.
+        value: Value tensor matching the shape, dtype, and device of your use case.
+    Returns:
+        (List[Tuple[int, int]]): List of tuples of two integers corresponding to query and KV tile
+            sizes in the *backward pass*.
+    """
+    if not can_run_cutlass_fmha(
+        query=query,
+        key=key,
+        value=value,
+        is_causal=False,
+        is_varlen=False,
+        raise_error=False,
+    ):
+        return []
+    return get_all_cutlass_fmha_backward_configs(
+        input_tensor=key if key.shape[-1] >= value.shape[-1] else value
+    )
+def get_configs_for_cutlass_fna(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+) -> List[CutlassFnaForwardConfigType]:
+    """Returns CUTLASS FNA configurations compatible with input tensors, if any.
+    Checks first if a CUDA tensor, and on a device with compute capability >= 5.0, and if so,
+    returns *forward pass* configurations compatible with the specific compute capability, tensor
+    dtype and head dim, and according to the rank of the token layout (1D/2D/3D).
+    Each configuration for this operation is a tuple of two integer tuples: `(q_tile_shape,
+    kv_tile_shape)`. These are arguments to [natten.na1d][natten.na1d], [natten.na2d][natten.na2d],
+    and [natten.na3d][natten.na3d].
+    Args:
+        query: Query tensor matching the shape, dtype, and device of your use case.
+        key:   Key tensor matching the shape, dtype, and device of your use case.
+        value: Value tensor matching the shape, dtype, and device of your use case.
+    Returns:
+        (List[Tuple[tuple, tuple]]): List of tuples of two integer tuples corresponding to query
+            and KV tile *shapes*.
+    """
+    if not can_run_cutlass_fna(query=query, key=key, value=value, raise_error=False):
+        return []
+    return get_all_cutlass_fna_forward_configs(input_tensor=query)
+def get_bwd_configs_for_cutlass_fna(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+) -> List[CutlassFnaBackwardConfigType]:
+    """Returns CUTLASS FNA backward pass configurations compatible with input tensors, if any.
+    Checks first if a CUDA tensor, and on a device with compute capability >= 5.0, and if so,
+    returns *backward pass* configurations compatible with the specific compute capability, tensor
+    dtype and head dim.
+    Each configuration for this operation is a tuple of two integers: `(backward_q_tile_shape,
+    backward_kv_tile_shape)`. These are arguments to [natten.na1d][natten.na1d],
+    [natten.na2d][natten.na2d], and [natten.na3d][natten.na3d].
+    Args:
+        query: Query tensor matching the shape, dtype, and device of your use case.
+        key:   Key tensor matching the shape, dtype, and device of your use case.
+        value: Value tensor matching the shape, dtype, and device of your use case.
+    Returns:
+        (List[Tuple[tuple, tuple]]): List of tuples of two integer tuples corresponding to query
+            and KV tile *shapes* in the *backward pass*.
+    """
+    if not can_run_cutlass_fna(query=query, key=key, value=value, raise_error=False):
+        return []
+    return get_all_cutlass_fna_backward_configs(
+        input_tensor=key if key.shape[-1] >= value.shape[-1] else value
+    )
+### Flex
+def get_configs_for_flex_fmha(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    torch_compile: bool = False,
+) -> List[FlexFmhaForwardConfigType]:
+    """Returns Flex FMHA configurations compatible with input tensors, if any.
+    Each configuration for this operation is a tuple of two integers: `(q_tile_size,
+    kv_tile_size)`. These are arguments to [natten.attention][natten.attention].
+    Not specifying these arguments while backend is Flex will default to `q_tile_size = 64` and
+    `kv_tile_size = 64`.
+    Args:
+        query: Query tensor matching the shape, dtype, and device of your use case.
+        key:   Key tensor matching the shape, dtype, and device of your use case.
+        value: Value tensor matching the shape, dtype, and device of your use case.
+        torch_compile: Whether or not you intend to use compiled block mask and flex attention kernel.
+    Returns:
+        (List[Tuple[int, int]]): List of tuples of two integers corresponding to query and KV tile
+            sizes.
+    """
+    if not can_run_flex_attention(
+        query=query,
+        key=key,
+        value=value,
+        torch_compile=torch_compile,
+        raise_error=False,
+    ):
+        return []
+    return get_all_flex_fmha_forward_configs(input_tensor=query)
+def get_configs_for_flex_fna(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    torch_compile: bool = False,
+) -> List[FlexFnaForwardConfigType]:
+    """Returns Flex FNA configurations compatible with input tensors, if any.
+    Each configuration for this operation is a tuple of two integer tuples: `(q_tile_shape,
+    kv_tile_shape)`. These are arguments to [natten.na1d][natten.na1d], [natten.na2d][natten.na2d],
+    and [natten.na3d][natten.na3d].
+    Not specifying these arguments while backend is Flex will default to single-dimensional tiling,
+    and will not use our Token Permutation approach. By explicitly specifying tile shapes, you will
+    automatically use our Token Permutation approach, which saves you the most compute.
+    Args:
+        query: Query tensor matching the shape, dtype, and device of your use case.
+        key:   Key tensor matching the shape, dtype, and device of your use case.
+        value: Value tensor matching the shape, dtype, and device of your use case.
+        torch_compile: Whether or not you intend to use compiled block mask and flex attention kernel.
+    Returns:
+        (List[Tuple[tuple, tuple]]): List of tuples of two integer tuples corresponding to query
+            and KV tile *shapes*.
+    """
+    if not can_run_flex_attention(
+        query=query,
+        key=key,
+        value=value,
+        torch_compile=torch_compile,
+        raise_error=False,
+    ):
+        return []
+    return get_all_flex_fna_forward_configs(input_tensor=query)

build/torch210-cu126-aarch64-linux/backends/configs/checks.py ADDED Viewed

	@@ -0,0 +1,750 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+import functools
+import math
+import torch
+from torch import Tensor
+from ..._environment import _IS_TORCH_COMPILE_SUPPORTED, _TORCH_VERSION
+from ..._libnatten import HAS_LIBNATTEN
+from ...context import is_flex_compile_allowed, is_flex_compile_backprop_allowed
+from ...utils.checks import fmha_tensor_checks, log_or_raise_error, na_tensor_checks
+from ...utils.device import get_device_cc, is_cpu, is_cuda, is_rocm
+from ...utils.dtype import is_fp8
+### Blackwell FMHA/FNA
+def can_run_cutlass_blackwell_fmha(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    is_causal: bool,
+    is_varlen: bool,
+    raise_error: bool = False,
+) -> bool:
+    target_fn = functools.partial(log_or_raise_error, raise_error=raise_error)
+    if not HAS_LIBNATTEN:
+        target_fn("Can't run Blackwell FMHA; NATTEN was not built with libnatten.")
+        return False
+    if not fmha_tensor_checks(
+        query,
+        key,
+        value,
+        must_match_head_dims=True,
+        supports_gqa_mqa=True,
+        raise_error=raise_error,
+        backend_name="Blackwell FMHA",
+    ):
+        return False
+    if query.dim() != 4:
+        target_fn(
+            f"Blackwell FMHA expects rank-4 input tensors, got {query.shape=}.",
+            exception=ValueError,
+        )
+        return False
+    if not is_cuda(query.device):
+        target_fn("Can't run Blackwell FMHA; not a CUDA tensor.")
+        return False
+    device_cc = get_device_cc(query.device)
+    if device_cc not in [100, 103]:
+        target_fn(
+            "Can't run Blackwell FMHA; tensor was on CUDA device with "
+            f"compute capability {device_cc}, expected 100 or 103."
+        )
+        return False
+    head_dim = query.shape[-1]
+    head_dim_v = value.shape[-1]
+    if head_dim != head_dim_v:
+        target_fn(
+            "Can't run Blackwell FMHA; it does not support different head dims for QK and V, "
+            f"got {head_dim=}, {head_dim_v=}.",
+            exception=ValueError,
+        )
+        return False
+    if query.dtype not in [
+        torch.float16,
+        torch.bfloat16,
+        torch.float8_e4m3fn,
+        torch.float8_e5m2,
+    ]:
+        target_fn(
+            "Can't run Blackwell FMHA; it only supports FP16, BF16, FP8-E4M3 and FP8-E5M2.",
+            exception=ValueError,
+        )
+        return False
+    requires_grad = query.requires_grad or key.requires_grad or value.requires_grad
+    if requires_grad and is_fp8(query.dtype):
+        target_fn(
+            "Blackwell FMHA does not support FP8 backward pass, but "
+            f"got {requires_grad=}, {query.dtype=}.",
+            exception=ValueError,
+        )
+        return False
+    if head_dim > 128:
+        target_fn(
+            f"Can't run Blackwell FMHA; maximum supported head dim is 128, got {head_dim}.",
+            exception=NotImplementedError,
+        )
+        return False
+    if is_fp8(query.dtype):
+        if head_dim < 16 or head_dim % 16 != 0:
+            target_fn(
+                "Can't run Blackwell FMHA; FP8 requires head dims that are multiples of 16 "
+                f"(minimum 16), got {head_dim}.",
+                exception=NotImplementedError,
+            )
+            return False
+    else:
+        if head_dim < 8 or head_dim % 8 != 0:
+            target_fn(
+                "Can't run Blackwell FMHA; FP16 and BF16 require head dims that are multiples "
+                f"of 8 (minimum 8), got {head_dim}.",
+                exception=NotImplementedError,
+            )
+            return False
+    return True
+def can_run_cutlass_blackwell_fna(
+    query: Tensor, key: Tensor, value: Tensor, raise_error: bool = False
+) -> bool:
+    target_fn = functools.partial(log_or_raise_error, raise_error=raise_error)
+    if not HAS_LIBNATTEN:
+        target_fn("Can't run Blackwell FNA; NATTEN was not built with libnatten.")
+        return False
+    if not na_tensor_checks(
+        query,
+        key,
+        value,
+        must_match_head_dims=True,
+        supports_gqa_mqa=True,
+        raise_error=raise_error,
+        backend_name="Blackwell FNA",
+    ):
+        return False
+    if query.dim() not in [4, 5, 6]:
+        target_fn(
+            "Blackwell FNA expects 4-D, 5-D, or 6-D tensors as inputs (corresponding to NA1D, NA2D, and NA3D), "
+            f"got {query.dim()=}.",
+            exception=ValueError,
+        )
+        return False
+    if not is_cuda(query.device):
+        target_fn("Can't run Blackwell FNA; not a CUDA tensor.")
+        return False
+    device_cc = get_device_cc(query.device)
+    if device_cc not in [100, 103]:
+        target_fn(
+            "Can't run Blackwell FNA; tensor was on CUDA device with "
+            f"compute capability {device_cc}, expected 100 or 103."
+        )
+        return False
+    requires_grad = query.requires_grad or key.requires_grad or value.requires_grad
+    if requires_grad and torch.are_deterministic_algorithms_enabled():
+        target_fn(
+            "Can't run Blackwell FNA; its backprop does not have a deterministic mode, but "
+            "PyTorch's deterministic mode was enabled.",
+            exception=NotImplementedError,
+        )
+        return False
+    head_dim = query.shape[-1]
+    head_dim_v = value.shape[-1]
+    if head_dim != head_dim_v:
+        target_fn(
+            "Can't run Blackwell FNA; it does not support different head dims for QK and V, "
+            f"got {head_dim=}, {head_dim_v=}.",
+            exception=ValueError,
+        )
+        return False
+    if query.dtype not in [
+        torch.float16,
+        torch.bfloat16,
+        torch.float8_e4m3fn,
+        torch.float8_e5m2,
+    ]:
+        target_fn(
+            "Can't run Blackwell FNA; it only supports FP16, BF16, FP8-E4M3 and FP8-E5M2.",
+            exception=ValueError,
+        )
+        return False
+    if requires_grad and is_fp8(query.dtype):
+        target_fn(
+            "Blackwell FNA does not support FP8 backward pass, but "
+            f"got {requires_grad=}, {query.dtype=}.",
+            exception=ValueError,
+        )
+        return False
+    if head_dim > 128:
+        target_fn(
+            f"Can't run Blackwell FNA; maximum supported head dim is 128, got {head_dim}.",
+            exception=NotImplementedError,
+        )
+        return False
+    if is_fp8(query.dtype):
+        if head_dim < 16 or head_dim % 16 != 0:
+            target_fn(
+                "Can't run Blackwell FNA; FP8 requires head dims that are multiples of 16 "
+                f"(minimum 16), got {head_dim}.",
+                exception=NotImplementedError,
+            )
+            return False
+    else:
+        if head_dim < 8 or head_dim % 8 != 0:
+            target_fn(
+                "Can't run Blackwell FNA; FP16 and BF16 require head dims that are multiples "
+                f"of 8 (minimum 8), got {head_dim}.",
+                exception=NotImplementedError,
+            )
+            return False
+    return True
+### Hopper FMHA/FNA
+def can_run_cutlass_hopper_fmha(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    is_causal: bool,
+    is_varlen: bool,
+    raise_error: bool = False,
+) -> bool:
+    target_fn = functools.partial(log_or_raise_error, raise_error=raise_error)
+    if not HAS_LIBNATTEN:
+        target_fn("Can't run Hopper FMHA; NATTEN was not built with libnatten.")
+        return False
+    if not fmha_tensor_checks(
+        query,
+        key,
+        value,
+        must_match_head_dims=True,
+        supports_gqa_mqa=True,  # NOTE: supports GQA in only by virtue of repeating heads manually
+        raise_error=raise_error,
+        backend_name="Hopper FMHA",
+    ):
+        return False
+    if query.dim() != 4:
+        target_fn(
+            f"Hopper FMHA expects rank-4 input tensors, got {query.shape=}.",
+            exception=ValueError,
+        )
+        return False
+    if not is_cuda(query.device):
+        target_fn("Can't run Hopper FMHA; not a CUDA tensor.")
+        return False
+    device_cc = get_device_cc(query.device)
+    if device_cc != 90:
+        target_fn(
+            "Can't run Hopper FMHA; tensor was on CUDA device with "
+            f"compute capability {device_cc}, expected 90."
+        )
+        return False
+    head_dim = query.shape[-1]
+    head_dim_v = value.shape[-1]
+    if head_dim != head_dim_v:
+        target_fn(
+            "Can't run Hopper FMHA; it does not support different head dims for QK and V, "
+            f"got {head_dim=}, {head_dim_v=}.",
+            exception=ValueError,
+        )
+        return False
+    requires_grad = query.requires_grad or key.requires_grad or value.requires_grad
+    if requires_grad and head_dim not in [32, 64, 128]:
+        target_fn(
+            f"Can't run Hopper FMHA; it does not support backpropagation for {head_dim=} yet; "
+            "only head dims 32, 64, and 128 are allowed.",
+            exception=NotImplementedError,
+        )
+        return False
+    if requires_grad and torch.are_deterministic_algorithms_enabled():
+        target_fn(
+            "Can't run Hopper FMHA; its backprop does not have a deterministic mode, but "
+            "PyTorch's deterministic mode was enabled.",
+            exception=NotImplementedError,
+        )
+        return False
+    if query.dtype not in [torch.float16, torch.bfloat16]:
+        target_fn(
+            "Can't run Hopper FMHA; it only supports FP16 and BF16 for now.",
+            exception=ValueError,
+        )
+        return False
+    if head_dim not in [32, 64, 128, 256]:
+        target_fn(
+            "Can't run Hopper FMHA; it only supports head dims 32, 64, 128, and 256 for now.",
+            exception=NotImplementedError,
+        )
+        return False
+    return True
+def can_run_cutlass_hopper_fna(
+    query: Tensor, key: Tensor, value: Tensor, raise_error: bool = False
+) -> bool:
+    target_fn = functools.partial(log_or_raise_error, raise_error=raise_error)
+    if not HAS_LIBNATTEN:
+        target_fn("Can't run Hopper FNA; NATTEN was not built with libnatten.")
+        return False
+    if not na_tensor_checks(
+        query,
+        key,
+        value,
+        must_match_head_dims=True,
+        supports_gqa_mqa=True,  # NOTE: supports GQA in only by virtue of repeating heads manually
+        raise_error=raise_error,
+        backend_name="Hopper FNA",
+    ):
+        return False
+    if query.dim() not in [4, 5, 6]:
+        target_fn(
+            "Hopper FNA expects 4-D, 5-D, or 6-D tensors as inputs (corresponding to NA1D, NA2D, and NA3D), "
+            f"got {query.dim()=}.",
+            exception=ValueError,
+        )
+        return False
+    if not is_cuda(query.device):
+        target_fn("Can't run Hopper FNA; not a CUDA tensor.")
+        return False
+    device_cc = get_device_cc(query.device)
+    if device_cc != 90:
+        target_fn(
+            "Can't run Hopper FNA; tensor was on CUDA device with "
+            f"compute capability {device_cc}, expected 90."
+        )
+        return False
+    head_dim = query.shape[-1]
+    head_dim_v = value.shape[-1]
+    if head_dim != head_dim_v:
+        target_fn(
+            "Can't run Hopper FNA; it does not support different head dims for QK and V, "
+            f"got {head_dim=}, {head_dim_v=}.",
+            exception=ValueError,
+        )
+        return False
+    requires_grad = query.requires_grad or key.requires_grad or value.requires_grad
+    if requires_grad and head_dim not in [32, 64, 128]:
+        target_fn(
+            f"Can't run Hopper FNA; it does not support backpropagation for {head_dim=} yet; "
+            "only head dims 32, 64, and 128 are allowed.",
+            exception=NotImplementedError,
+        )
+        return False
+    if requires_grad and torch.are_deterministic_algorithms_enabled():
+        target_fn(
+            "Can't run Hopper FNA; its backprop does not have a deterministic mode, but "
+            "PyTorch's deterministic mode was enabled.",
+            exception=NotImplementedError,
+        )
+        return False
+    if query.dtype not in [torch.float16, torch.bfloat16]:
+        target_fn(
+            "Can't run Hopper FNA; it only supports FP16 and BF16 for now.",
+            exception=ValueError,
+        )
+        return False
+    if head_dim not in [32, 64, 128, 256]:
+        target_fn(
+            "Can't run Hopper FNA; it only supports head dims 32, 64, 128, and 256 for now.",
+            exception=NotImplementedError,
+        )
+        return False
+    return True
+### CUTLASS FMHA/FNA
+def can_run_cutlass_fmha(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    is_causal: bool,
+    is_varlen: bool,
+    raise_error: bool = False,
+) -> bool:
+    target_fn = functools.partial(log_or_raise_error, raise_error=raise_error)
+    if not HAS_LIBNATTEN:
+        target_fn("Can't run CUTLASS FMHA; NATTEN was not built with libnatten.")
+        return False
+    if not fmha_tensor_checks(
+        query,
+        key,
+        value,
+        must_match_head_dims=False,
+        supports_gqa_mqa=True,  # NOTE: supports GQA in only by virtue of repeating heads manually
+        raise_error=raise_error,
+        backend_name="CUTLASS FMHA",
+    ):
+        return False
+    if query.dim() != 4:
+        target_fn(
+            f"FMHA expects rank-4 input tensors, got {query.shape=}.",
+            exception=ValueError,
+        )
+        return False
+    if not is_cuda(query.device):
+        target_fn("Can't run CUTLASS FMHA; not a CUDA tensor.")
+        return False
+    device_cc = get_device_cc(query.device)
+    if device_cc < 60:
+        target_fn(
+            "CUTLASS FMHA only supports CUDA devices with compute capability 60 or higher, "
+            f"got {device_cc}."
+        )
+        return False
+    head_dim = query.shape[-1]
+    head_dim_v = value.shape[-1]
+    if query.dtype not in [torch.float32, torch.float16, torch.bfloat16]:
+        target_fn(
+            "Can't run CUTLASS FMHA; it only supports FP32, FP16, and BF16.",
+            exception=ValueError,
+        )
+        return False
+    if head_dim % 8 != 0:
+        target_fn(
+            "Can't run CUTLASS FMHA; it only supports head dims that are multiples of 8.",
+            exception=ValueError,
+        )
+        return False
+    if head_dim_v % 8 != 0:
+        target_fn(
+            "Can't run CUTLASS FMHA; it only supports head dims that are multiples of 8, "
+            f"got {head_dim_v=}.",
+            exception=ValueError,
+        )
+        return False
+    if max(head_dim, head_dim_v) > 2**16:
+        target_fn(
+            f"Can't run CUTLASS FMHA; it supports max head dim of {2**16}, "
+            f"got {head_dim=}, {head_dim_v=}.",
+            exception=ValueError,
+        )
+        return False
+    return True
+def can_run_cutlass_fna(
+    query: Tensor, key: Tensor, value: Tensor, raise_error: bool = False
+) -> bool:
+    target_fn = functools.partial(log_or_raise_error, raise_error=raise_error)
+    if not HAS_LIBNATTEN:
+        target_fn("Can't run CUTLASS FNA; NATTEN was not built with libnatten.")
+        return False
+    if not na_tensor_checks(
+        query,
+        key,
+        value,
+        must_match_head_dims=False,
+        supports_gqa_mqa=True,  # NOTE: supports GQA in only by virtue of repeating heads manually
+        raise_error=raise_error,
+        backend_name="CUTLASS FNA",
+    ):
+        return False
+    if query.dim() not in [4, 5, 6]:
+        target_fn(
+            "CUTLASS FNA expects 4-D, 5-D, or 6-D tensors as inputs (corresponding to NA1D, NA2D, and NA3D), "
+            f"got {query.dim()=}.",
+            exception=ValueError,
+        )
+        return False
+    if not is_cuda(query.device):
+        target_fn("Can't run CUTLASS FNA; not a CUDA tensor.")
+        return False
+    device_cc = get_device_cc(query.device)
+    if device_cc < 60:
+        target_fn(
+            "CUTLASS FNA only supports CUDA devices with compute capability 60 or higher, "
+            f"got {device_cc}."
+        )
+        return False
+    head_dim = query.shape[-1]
+    head_dim_v = value.shape[-1]
+    if query.dtype not in [torch.float32, torch.float16, torch.bfloat16]:
+        target_fn(
+            "Can't run CUTLASS FNA; it only supports FP32, FP16, and BF16.",
+            exception=ValueError,
+        )
+        return False
+    if head_dim % 8 != 0:
+        target_fn(
+            "Can't run CUTLASS FNA; it only supports head dims that are multiples of 8, "
+            f"got {head_dim=}.",
+            exception=ValueError,
+        )
+        return False
+    if head_dim_v % 8 != 0:
+        target_fn(
+            "Can't run CUTLASS FNA; it only supports head dims that are multiples of 8, "
+            f"got {head_dim_v=}.",
+            exception=ValueError,
+        )
+        return False
+    if max(head_dim, head_dim_v) > 2**16:
+        target_fn(
+            f"Can't run CUTLASS FNA; it supports max head dim of {2**16}, "
+            f"got {head_dim=}, {head_dim_v=}.",
+            exception=ValueError,
+        )
+        return False
+    return True
+### Flex FMHA/FNA
+_FLEX_SUPPORTED = _TORCH_VERSION >= [2, 7]
+_FLEX_COMPILE_SUPPORTED = _TORCH_VERSION >= [2, 7] and _IS_TORCH_COMPILE_SUPPORTED
+def can_run_flex_attention(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    torch_compile: bool,
+    is_causal: bool = False,
+    is_varlen: bool = False,
+    raise_error: bool = False,
+) -> bool:
+    target_fn = functools.partial(log_or_raise_error, raise_error=raise_error)
+    if is_causal:
+        target_fn("Flex FMHA doesn't support causal mask yet.")
+        return False
+    if is_varlen:
+        target_fn("Flex FMHA doesn't support variable length inputs (varlen).")
+        return False
+    if not _FLEX_SUPPORTED:
+        target_fn("Can't run NATTEN with Flex Attention with torch < 2.7.")
+        return False
+    if torch_compile and not _FLEX_COMPILE_SUPPORTED:
+        target_fn("Can't run NATTEN with Flex Attention (compiled).)")
+        return False
+    if torch_compile and not is_flex_compile_allowed():
+        target_fn(
+            "NATTEN does not allow compiling Flex Attention. This is because we cannot verify "
+            "Flex's correctness in all scenarios through NATTEN's tests. You can choose to override "
+            "this, though it is discouraged, as it may affect your results significantly, "
+            "by doing:\n"
+            "  from ... import allow_flex_compile\n"
+            "  allow_flex_compile()\n"
+        )
+        return False
+    requires_grad = query.requires_grad or key.requires_grad or value.requires_grad
+    if torch_compile and requires_grad and not is_flex_compile_backprop_allowed():
+        target_fn(
+            "NATTEN does not allow compiling Flex Attention for backpropagation "
+            "({q,k,v}.requires_grad=True). This is because we cannot verify Flex's correctness "
+            "in all scenarios through NATTEN's tests. You can choose to override this, though "
+            "it is HIGHLY discouraged, as it may affect the results of your training significantly, "
+            "by doing:\n"
+            "  from ... import allow_flex_compile_backprop\n"
+            "  allow_flex_compile_backprop()\n"
+        )
+        return False
+    # TODO: can we just have different checks for FMHA vs FNA, like the rest of the backends?
+    if query.dim() == 4 and key.dim() == 4 and query.shape[1] != key.shape[1]:
+        supported = fmha_tensor_checks(
+            query,
+            key,
+            value,
+            must_match_head_dims=True,
+            supports_gqa_mqa=True,
+            raise_error=raise_error,
+            backend_name="Flex FMHA",
+        )
+    else:
+        supported = na_tensor_checks(
+            query,
+            key,
+            value,
+            must_match_head_dims=True,
+            supports_gqa_mqa=True,
+            raise_error=raise_error,
+            backend_name="Flex FMHA/FNA",
+        )
+    if not supported:
+        return False
+    if query.dim() not in [4, 5, 6]:
+        target_fn(
+            "Flex backend expects 4-D, 5-D, or 6-D tensors as inputs (corresponding to FMHA/NA1D, "
+            f"NA2D, and NA3D), got {query.dim()=}.",
+            exception=ValueError,
+        )
+        return False
+    if not is_cuda(query.device):
+        if not is_cpu(query.device) and not is_rocm(query.device):
+            target_fn(
+                "Can't run Flex Attention; tensor is not on a CUDA, ROCm, or CPU device: "
+                f"{query.device.type}"
+            )
+            return False
+        # TODO: check if ROCm device supports torch.compile/triton?
+    else:
+        device_cc = get_device_cc(query.device)
+        if device_cc < 70:
+            target_fn(
+                "Flex Attention (compiled) only supports CUDA devices with compute capability "
+                f"70 or higher, got {device_cc}."
+            )
+            return False
+    head_dim = query.shape[-1]
+    head_dim_v = value.shape[-1]
+    if head_dim != head_dim_v:
+        target_fn(
+            "Can't run NATTEN with Flex Attention; we don't support different head dims for QK and "
+            f"V in this backend yet, got {head_dim=}, {head_dim_v=}.",
+            exception=ValueError,
+        )
+        return False
+    if not torch_compile and query.dtype not in [
+        torch.float32,
+        torch.float16,
+        torch.bfloat16,
+    ]:
+        target_fn(
+            "Can't run NATTEN with Flex Attention; we only support FP32, FP16, and BF16 for now.",
+            exception=ValueError,
+        )
+        return False
+    if torch_compile and query.dtype not in [torch.float16, torch.bfloat16]:
+        target_fn(
+            "Can't run NATTEN with Flex Attention (compiled); we only support FP32, FP16, and BF16 for now.",
+            exception=ValueError,
+        )
+        return False
+    if torch_compile and (
+        head_dim < 32 or head_dim > 512 or not math.log2(head_dim).is_integer()
+    ):
+        target_fn(
+            "Can't run NATTEN with Flex Attention (compiled); we only allow 32 <= head_dim <= 512 "
+            f"and only powers of two, got {head_dim}.",
+            exception=ValueError,
+        )
+        return False
+    if not torch_compile and (
+        head_dim < 8 or head_dim > 512 or not math.log2(head_dim).is_integer()
+    ):
+        target_fn(
+            "Can't run NATTEN with Flex Attention (not compiled); we only allow 8 <= head_dim <= 512 "
+            f"and only powers of two, got {head_dim}.",
+            exception=ValueError,
+        )
+        return False
+    return True

build/torch210-cu126-aarch64-linux/backends/configs/cutlass/__init__.py ADDED Viewed

	@@ -0,0 +1,422 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+from typing import List, Optional, Tuple
+import torch
+from torch import Tensor
+# FNA/FMHA forward supports 64x64 and 32x128 GEMM configs in all
+# use cases. Some architectures (SM80 and SM90 )have more shared
+# memory so they can handle 64x128 GEMMs.
+from ....backends.configs.cutlass.fna_backward_128x128 import (
+    _FNA_BACKWARD_128x128_TILE_SIZES,
+)
+from ....backends.configs.cutlass.fna_backward_128x64 import (
+    _FNA_BACKWARD_128x64_TILE_SIZES,
+)
+from ....backends.configs.cutlass.fna_backward_64x64 import (
+    _FNA_BACKWARD_64x64_TILE_SIZES,
+)
+# FNA/FMHA backward supports 64x64 GEMM configs in all
+# use cases. Some architectures have more shared memory
+# so they can handle 128x64 or 128x128 GEMMs, but that
+# is also dependent on the GEMM K.
+from ....backends.configs.cutlass.fna_forward_32x128 import (
+    _FNA_FORWARD_32x128_TILE_SIZES,
+)
+from ....backends.configs.cutlass.fna_forward_64x128 import (
+    _FNA_FORWARD_64x128_TILE_SIZES,
+)
+from ....backends.configs.cutlass.fna_forward_64x64 import (
+    _FNA_FORWARD_64x64_TILE_SIZES,
+)
+from ...._types import (
+    CutlassFmhaBackwardConfigType,
+    CutlassFmhaForwardConfigType,
+    CutlassFnaBackwardConfigType,
+    CutlassFnaForwardConfigType,
+    DimensionType,
+)
+from ....utils.checks import check_tile_shape
+from ....utils.device import get_device_cc, is_cuda
+def _get_default_tile_shapes_forward(
+    na_dim: int,
+) -> CutlassFnaForwardConfigType:
+    assert na_dim in [1, 2, 3]
+    if na_dim == 1:
+        return ((64,), (64,))
+    if na_dim == 2:
+        return ((8, 8), (8, 8))
+    if na_dim == 3:
+        return ((4, 4, 4), (4, 4, 4))
+    raise NotImplementedError()
+def get_all_tile_shapes_forward(
+    input_tensor: Tensor,
+) -> List[CutlassFnaForwardConfigType]:
+    assert input_tensor.dim() in [4, 5, 6]
+    na_dim = input_tensor.dim() - 3  # batch, heads, head_dim
+    device = input_tensor.device
+    if not is_cuda(device):
+        return []
+    # DC-class GPUs have more shared memory
+    if get_device_cc(device) in [80, 90, 100, 103]:
+        return (
+            _FNA_FORWARD_32x128_TILE_SIZES[na_dim]
+            + _FNA_FORWARD_64x64_TILE_SIZES[na_dim]
+            + _FNA_FORWARD_64x128_TILE_SIZES[na_dim]
+        )
+    return (
+        _FNA_FORWARD_32x128_TILE_SIZES[na_dim] + _FNA_FORWARD_64x64_TILE_SIZES[na_dim]
+    )
+# For FMHA
+def get_all_tile_sizes_forward(
+    input_tensor: Tensor,
+) -> List[CutlassFmhaForwardConfigType]:
+    if input_tensor.dim() != 4:
+        raise ValueError("Only 4-D tensors are supported in FMHA.")
+    tile_shapes = get_all_tile_shapes_forward(input_tensor)
+    assert all(len(q_t) == len(kv_t) == 1 for q_t, kv_t in tile_shapes)
+    tile_sizes = [(q_t[0], kv_t[0]) for q_t, kv_t in tile_shapes]
+    return tile_sizes
+def get_default_forward_config(
+    input_tensor: Tensor, dilation: Optional[DimensionType] = None
+) -> CutlassFnaForwardConfigType:
+    assert input_tensor.dim() in [4, 5, 6]
+    na_dim = input_tensor.dim() - 3  # batch, heads, head_dim
+    return _get_default_tile_shapes_forward(na_dim)
+get_all_forward_configs = get_all_tile_shapes_forward
+get_all_fmha_forward_configs = get_all_tile_sizes_forward
+def check_cutlass_fna_forward_config(
+    input_tensor: Tensor,
+    dilation: Optional[DimensionType] = None,
+    q_tile_shape: Optional[DimensionType] = None,
+    kv_tile_shape: Optional[DimensionType] = None,
+) -> CutlassFnaForwardConfigType:
+    assert input_tensor.dim() in [4, 5, 6]
+    na_dim = input_tensor.dim() - 3  # batch, heads, head_dim
+    if (q_tile_shape is None) ^ (kv_tile_shape is None):
+        raise ValueError(
+            "Please specify both q_tile_shape and kv_tile_shape, or neither one. "
+            f"Got {q_tile_shape=}, {kv_tile_shape=}."
+        )
+    if q_tile_shape is None and kv_tile_shape is None:
+        return get_default_forward_config(input_tensor=input_tensor, dilation=dilation)
+    q_tile_shape = check_tile_shape(q_tile_shape)
+    kv_tile_shape = check_tile_shape(kv_tile_shape)
+    tile_shapes = get_all_tile_shapes_forward(input_tensor=input_tensor)
+    for q_t, kv_t in tile_shapes:
+        if q_t == q_tile_shape and kv_t == kv_tile_shape:
+            return (q_t, kv_t)  # type: ignore
+    # Fail and make suggestions
+    device_cc = get_device_cc(input_tensor.device)
+    MAX_EXAMPLES = 3
+    examples = ""
+    for i, (q_t, kv_t) in enumerate(tile_shapes):
+        examples += f"\n  q_tile_shape={q_t}, kv_tile_shape={kv_t}"
+        if i > MAX_EXAMPLES:
+            break
+    raise ValueError(
+        f"Invalid configuration for CUTLASS FNA-{na_dim}D. "
+        f"Q tile shape {q_tile_shape} and KV tile shape {kv_tile_shape} "
+        f"are not among the {len(tile_shapes)} configurations implementable "
+        f"with CUTLASS 2.X FNA for SM{device_cc}. "
+        "Try selecting a combination from: \n"
+        "  natten.get_configs_for_cutlass_fna(q, k, v)"
+        "\n"
+        "Here's a few examples of available combinations for your use case:\n"
+        f"{examples}"
+    )
+def check_cutlass_fmha_forward_config(
+    input_tensor: Tensor,
+    q_tile_size: Optional[int] = None,
+    kv_tile_size: Optional[int] = None,
+) -> CutlassFmhaForwardConfigType:
+    assert input_tensor.dim() == 4
+    if (q_tile_size is None) ^ (kv_tile_size is None):
+        raise ValueError(
+            "Please specify both q_tile_size and kv_tile_size, or neither one. "
+            f"Got {q_tile_size=}, {kv_tile_size=}."
+        )
+    if q_tile_size is None and kv_tile_size is None:
+        q_tile_shape, kv_tile_shape = get_default_forward_config(
+            input_tensor=input_tensor
+        )
+        assert len(q_tile_shape) == len(kv_tile_shape) == 1
+        return (q_tile_shape[0], kv_tile_shape[0])
+    tile_sizes = get_all_tile_sizes_forward(input_tensor=input_tensor)
+    for q_t, kv_t in tile_sizes:
+        if q_t == q_tile_size and kv_t == kv_tile_size:
+            return (q_t, kv_t)
+    # Fail and make suggestions
+    device_cc = get_device_cc(input_tensor.device)
+    MAX_EXAMPLES = 3
+    examples = ""
+    for i, (q_t, kv_t) in enumerate(tile_sizes):
+        examples += f"\n  q_tile_size={q_t}, kv_tile_size={kv_t}"
+        if i > MAX_EXAMPLES:
+            break
+    raise ValueError(
+        f"Invalid configuration for CUTLASS FMHA. "
+        f"Q tile size {q_tile_size} and KV tile size {kv_tile_size} "
+        f"are not among the {len(tile_sizes)} configurations implementable "
+        f"with CUTLASS 2.X FMHA for SM{device_cc}. "
+        "Try selecting a combination from: \n"
+        "  natten.get_configs_for_cutlass_fmha(q, k, v)"
+        "\n"
+        "Here's a few examples of available combinations for your use case:\n"
+        f"{examples}"
+    )
+###### Backward
+def _get_default_tile_shapes_backward(
+    na_dim: int,
+) -> Tuple[DimensionType, DimensionType]:
+    assert na_dim in [1, 2, 3]
+    if na_dim == 1:
+        return ((64,), (64,))
+    if na_dim == 2:
+        return ((8, 8), (8, 8))
+    if na_dim == 3:
+        return ((4, 4, 4), (4, 4, 4))
+    raise NotImplementedError()
+def get_all_tile_shapes_backward(
+    input_tensor: Tensor,
+) -> List[CutlassFnaForwardConfigType]:
+    assert input_tensor.dim() in [4, 5, 6]
+    na_dim = input_tensor.dim() - 3  # batch, heads, head_dim
+    device = input_tensor.device
+    dtype = input_tensor.dtype
+    dim_per_head = input_tensor.shape[-1]
+    if not is_cuda(device):
+        return []
+    compute_cap = get_device_cc(device)
+    assert dtype in [torch.float32, torch.float16, torch.bfloat16]
+    # DC-class cards have extra shmem which allows larger tile sizes
+    dc_class_arches = [80, 90, 100, 103]
+    if dtype == torch.float32 and compute_cap not in dc_class_arches:
+        return _FNA_BACKWARD_64x64_TILE_SIZES[na_dim]
+    elif dtype == torch.float32:
+        return (
+            _FNA_BACKWARD_64x64_TILE_SIZES[na_dim]
+            + _FNA_BACKWARD_128x64_TILE_SIZES[na_dim]
+        )
+    if compute_cap == 70:
+        return (
+            _FNA_BACKWARD_64x64_TILE_SIZES[na_dim]
+            + _FNA_BACKWARD_128x64_TILE_SIZES[na_dim]
+        )
+    if compute_cap in dc_class_arches and dim_per_head <= 128:
+        return (
+            _FNA_BACKWARD_64x64_TILE_SIZES[na_dim]
+            + _FNA_BACKWARD_128x128_TILE_SIZES[na_dim]
+        )
+    elif compute_cap in dc_class_arches:
+        return (
+            _FNA_BACKWARD_64x64_TILE_SIZES[na_dim]
+            + _FNA_BACKWARD_128x64_TILE_SIZES[na_dim]
+        )
+    return _FNA_BACKWARD_64x64_TILE_SIZES[na_dim]
+# For FMHA
+def get_all_tile_sizes_backward(
+    input_tensor: Tensor,
+) -> List[CutlassFmhaForwardConfigType]:
+    if input_tensor.dim() != 4:
+        raise ValueError("Only 4-D tensors are supported in FMHA.")
+    tile_shapes = get_all_tile_shapes_backward(input_tensor)
+    assert all(len(q_t) == len(kv_t) == 1 for q_t, kv_t in tile_shapes)
+    tile_sizes = [(q_t[0], kv_t[0]) for q_t, kv_t in tile_shapes]
+    return tile_sizes
+def get_default_backward_config(
+    input_tensor: Tensor,
+) -> CutlassFnaBackwardConfigType:
+    assert input_tensor.dim() in [4, 5, 6]
+    na_dim = input_tensor.dim() - 3  # batch, heads, head_dim
+    q_tile_shape, kv_tile_shape = _get_default_tile_shapes_backward(na_dim)
+    return (q_tile_shape, kv_tile_shape)  # type: ignore
+get_all_backward_configs = get_all_tile_shapes_backward
+get_all_fmha_backward_configs = get_all_tile_sizes_backward
+def check_cutlass_fna_backward_config(
+    input_tensor: Tensor,
+    q_tile_shape: Optional[DimensionType] = None,
+    kv_tile_shape: Optional[DimensionType] = None,
+) -> CutlassFnaBackwardConfigType:
+    assert input_tensor.dim() in [4, 5, 6]
+    na_dim = input_tensor.dim() - 3  # batch, heads, head_dim
+    if (q_tile_shape is None) ^ (kv_tile_shape is None):
+        raise ValueError(
+            "Please specify both q_tile_shape and kv_tile_shape, or neither one. "
+            f"Got {q_tile_shape=}, {kv_tile_shape=}."
+        )
+    if q_tile_shape is None and kv_tile_shape is None:
+        return get_default_backward_config(
+            input_tensor=input_tensor,
+        )
+    q_tile_shape = check_tile_shape(q_tile_shape)
+    kv_tile_shape = check_tile_shape(kv_tile_shape)
+    tile_shapes = get_all_tile_shapes_backward(input_tensor)
+    for q_t, kv_t in tile_shapes:
+        if q_t == q_tile_shape and kv_t == kv_tile_shape:
+            return (q_t, kv_t)  # type: ignore
+    # Fail and make suggestions
+    device_cc = get_device_cc(input_tensor.device)
+    MAX_EXAMPLES = 3
+    examples = ""
+    for i, (q_t, kv_t) in enumerate(tile_shapes):
+        examples += f"\n  q_tile_shape={q_t}, kv_tile_shape={kv_t}"
+        if i > MAX_EXAMPLES:
+            break
+    raise ValueError(
+        f"Invalid configuration for CUTLASS FNA-{na_dim}D. "
+        f"Q tile shape {q_tile_shape} and KV tile shape {kv_tile_shape} "
+        f"are not among the {len(tile_shapes)} configurations implementable "
+        f"with CUTLASS 2.X FNA for SM{device_cc}, with input tensor shape "
+        f"{input_tensor.shape}. Try selecting a combination from: \n"
+        "  natten.get_bwd_configs_for_cutlass_fna(q, k, v)"
+        "\n"
+        "Here's a few examples of available combinations for your use case:\n"
+        f"{examples}"
+    )
+def check_cutlass_fmha_backward_config(
+    input_tensor: Tensor,
+    q_tile_size: Optional[int] = None,
+    kv_tile_size: Optional[int] = None,
+) -> CutlassFmhaBackwardConfigType:
+    assert input_tensor.dim() == 4
+    if (q_tile_size is None) ^ (kv_tile_size is None):
+        raise ValueError(
+            "Please specify both q_tile_size and kv_tile_size, or neither one. "
+            f"Got {q_tile_size=}, {kv_tile_size=}."
+        )
+    if q_tile_size is None and kv_tile_size is None:
+        q_tile_shape, kv_tile_shape = get_default_backward_config(
+            input_tensor=input_tensor
+        )
+        assert len(q_tile_shape) == len(kv_tile_shape) == 1
+        q_tile_size, kv_tile_size = q_tile_shape[0], kv_tile_shape[0]
+    tile_sizes = get_all_tile_sizes_backward(input_tensor)
+    for q_t, kv_t in tile_sizes:
+        if q_t == q_tile_size and kv_t == kv_tile_size:
+            return (q_t, kv_t)
+    # Fail and make suggestions
+    device_cc = get_device_cc(input_tensor.device)
+    MAX_EXAMPLES = 3
+    examples = ""
+    for i, (q_t, kv_t) in enumerate(tile_sizes):
+        examples += f"\n  q_tile_size={q_t}, kv_tile_size={kv_t}"
+        if i > MAX_EXAMPLES:
+            break
+    raise ValueError(
+        f"Invalid configuration for CUTLASS FMHA. "
+        f"Q tile size {q_tile_size} and KV tile size {kv_tile_size} "
+        f"are not among the {len(tile_sizes)} configurations implementable "
+        f"with CUTLASS 2.X FNA for SM{device_cc}, with input tensor shape "
+        f"{input_tensor.shape}. Try selecting a combination from: \n"
+        "  natten.get_bwd_configs_for_cutlass_fmha(q, k, v)"
+        "\n"
+        "Here's a few examples of available combinations for your use case:\n"
+        f"{examples}"
+    )

build/torch210-cu126-aarch64-linux/backends/configs/cutlass/backward_knobs.py ADDED Viewed

	@@ -0,0 +1,228 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+# Backward pass "knobs" for CUTLASS FNA/FMHA:
+#   - kv_splits: number of KV splits for parallelism
+#   - use_pt_reduction: whether to use PyTorch for delta computation
+#
+# These are independent of tile shape selection and are validated/defaulted
+# in the torch ops (torch_wrappers.py), not in the config selection logic.
+import itertools
+import math
+from typing import Optional
+from torch import Tensor
+from ....context import (
+    is_kv_parallelism_in_fused_na_enabled,
+    is_memory_usage_strict,
+    is_memory_usage_unrestricted,
+)
+from ...._types import DimensionType
+from ....utils.checks import check_dilation_arg, check_input_size_arg
+from ....utils.tuples import ceil_div_int, ceil_div_tuple
+def _get_max_grid_size_allowed() -> int:
+    if is_memory_usage_unrestricted():
+        return 65535
+    if is_memory_usage_strict():
+        return 1024
+    return 4096
+def get_min_splits(na_dim: int) -> DimensionType:
+    assert na_dim in [1, 2, 3]
+    return tuple(1 for _ in range(na_dim))  # type: ignore
+def get_max_splits(
+    input_shape: DimensionType, dilation: DimensionType, kv_tile_shape: DimensionType
+) -> DimensionType:
+    extent_per_dilation_group = ceil_div_tuple(input_shape, dilation)
+    return tuple(
+        ceil_div_int(x, t) for x, t in zip(extent_per_dilation_group, kv_tile_shape)
+    )  # type: ignore
+def _reduce_max_kv_splits(
+    na_dim: int,
+    kv_splits: DimensionType,
+    max_splits: int,
+) -> DimensionType:
+    assert isinstance(kv_splits, tuple)
+    assert na_dim in [1, 2, 3]
+    if na_dim == 1:
+        assert len(kv_splits) == 1
+        return (min(kv_splits[0], max_splits),)
+    if na_dim == 2:
+        assert len(kv_splits) == 2
+        splits_x = max(min(max_splits // 2, kv_splits[0]), 1)
+        splits_y = max(min(max_splits // splits_x, kv_splits[1]), 1)
+        assert (
+            0 < splits_x * splits_y <= max_splits
+        ), f"{splits_x=} * {splits_y=} does not fall in range [0, {max_splits}]"
+        return (splits_x, splits_y)
+    if na_dim == 3:
+        assert len(kv_splits) == 3
+        splits_x = max(min(max_splits // 3, kv_splits[0]), 1)
+        splits_y = max(min(max_splits // splits_x, kv_splits[1]), 1)
+        splits_z = max(min(max_splits // (splits_x * splits_y), kv_splits[2]), 1)
+        assert (
+            0 < splits_x * splits_y * splits_z <= max_splits
+        ), f"{splits_x=} * {splits_y=} * {splits_z=} does not fall in range [0, {max_splits}]"
+        return (splits_x, splits_y, splits_z)
+    raise NotImplementedError()
+def _get_possible_kv_splits(
+    min_splits: DimensionType,
+    max_splits: DimensionType,
+):
+    assert 0 < len(min_splits) == len(max_splits) < 4
+    na_dim = len(max_splits)
+    if na_dim == 1:
+        return itertools.product(
+            range(min_splits[0], max_splits[0] + 1),
+        )
+    if na_dim == 2:
+        assert len(min_splits) == len(max_splits) == 2
+        return itertools.product(
+            range(min_splits[0], max_splits[0] + 1),
+            range(min_splits[1], max_splits[1] + 1),
+        )
+    if na_dim == 3:
+        assert len(min_splits) == len(max_splits) == 3
+        return itertools.product(
+            range(min_splits[0], max_splits[0] + 1),
+            range(min_splits[1], max_splits[1] + 1),
+            range(min_splits[2], max_splits[2] + 1),
+        )
+    raise NotImplementedError()
+def get_default_kv_splits_backward(
+    input_tensor: Tensor,
+    kv_tile_shape: DimensionType,
+    deterministic: bool,
+    dilation: Optional[DimensionType] = None,
+    max_seqlen: Optional[DimensionType] = None,
+) -> DimensionType:
+    assert input_tensor.dim() in [4, 5, 6]
+    na_dim = input_tensor.dim() - 3  # batch, heads, head_dim
+    dilation = check_dilation_arg(na_dim, dilation)
+    input_shape: DimensionType = tuple(int(x) for x in input_tensor.shape[1 : na_dim + 1])  # type: ignore
+    if max_seqlen is not None:
+        input_shape = check_input_size_arg(na_dim, max_seqlen)
+    assert na_dim in [1, 2, 3]
+    if na_dim == 1:
+        kv_splits: DimensionType = (1,)
+    elif na_dim == 2:
+        kv_splits = (1, 1)
+    elif na_dim == 3:
+        kv_splits = (1, 1, 1)
+    if is_kv_parallelism_in_fused_na_enabled() and not deterministic:
+        kv_splits = get_max_splits(
+            input_shape, dilation=dilation, kv_tile_shape=kv_tile_shape
+        )
+        total_kv_splits = math.prod(kv_splits)
+        batch_size = input_tensor.shape[0]
+        num_heads = input_tensor.shape[-2]
+        num_dilation_splits = math.prod(dilation)
+        max_kv_splits_allowed = max(
+            1,
+            _get_max_grid_size_allowed()
+            // (batch_size * num_heads * num_dilation_splits),
+        )
+        if total_kv_splits > max_kv_splits_allowed:
+            kv_splits = _reduce_max_kv_splits(
+                na_dim=na_dim, kv_splits=kv_splits, max_splits=max_kv_splits_allowed
+            )
+    return kv_splits
+def check_fmha_kv_splits(
+    kv_splits: Optional[int],
+    input_tensor: Tensor,
+    kv_tile_size: int,
+    deterministic: bool,
+    max_seqlen: Optional[int] = None,
+) -> int:
+    if kv_splits is not None and isinstance(kv_splits, int):
+        seqlen_kv = input_tensor.shape[1] if max_seqlen is None else max_seqlen
+        num_kv_tiles = (seqlen_kv + kv_tile_size - 1) // kv_tile_size
+        assert num_kv_tiles > 0
+        return min(num_kv_tiles, kv_splits)
+    if kv_splits is None:
+        max_seqlen_tuple = None if max_seqlen is None else (max_seqlen,)
+        default_kv_splits: DimensionType = get_default_kv_splits_backward(
+            input_tensor=input_tensor,
+            deterministic=deterministic,
+            kv_tile_shape=(kv_tile_size,),
+            max_seqlen=max_seqlen_tuple,
+        )
+        assert len(default_kv_splits) == 1
+        return default_kv_splits[0]
+    raise ValueError(f"Invalid type {type(kv_splits)} for kv_splits.")
+def check_fna_kv_splits(
+    kv_splits: Optional[DimensionType],
+    input_tensor: Tensor,
+    kv_tile_shape: DimensionType,
+    deterministic: bool,
+    dilation: Optional[DimensionType] = None,
+) -> DimensionType:
+    if kv_splits is not None and isinstance(kv_splits, tuple):
+        na_dim = input_tensor.dim() - 3
+        dilation = check_dilation_arg(na_dim, dilation)
+        input_shape: DimensionType = tuple(int(x) for x in input_tensor.shape[1 : na_dim + 1])  # type: ignore
+        max_kv_splits = get_max_splits(
+            input_shape, dilation=dilation, kv_tile_shape=kv_tile_shape
+        )
+        return tuple(min(s, m) for s, m in zip(kv_splits, max_kv_splits))  # type: ignore
+    if kv_splits is None:
+        return get_default_kv_splits_backward(
+            deterministic=deterministic,
+            input_tensor=input_tensor,
+            kv_tile_shape=kv_tile_shape,
+            dilation=dilation,
+        )
+    raise ValueError(f"Invalid type {type(kv_splits)} for kv_splits.")

build/torch210-cu126-aarch64-linux/backends/configs/cutlass/fna_backward_128x128.py ADDED Viewed

	@@ -0,0 +1,304 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+from typing import Dict, List
+from ...._types import QKTileShapeType
+# NOTE: we're excluding tile shapes that include 1 just to
+# reduce the giant number of configs down to a reasonable
+# amount; otherwise autotuning would take more than a few
+# seconds per call which is unacceptable. Tile shapes with
+# 1s are rarely selected.
+_FNA_BACKWARD_128x128_TILE_SIZES: Dict[int, List[QKTileShapeType]] = {
+    1: [
+        ((128,), (128,)),
+    ],
+    2: [
+        ((64, 2), (64, 2)),
+        ((64, 2), (32, 4)),
+        ((64, 2), (16, 8)),
+        ((64, 2), (8, 16)),
+        ((64, 2), (4, 32)),
+        ((64, 2), (2, 64)),
+        ((32, 4), (64, 2)),
+        ((32, 4), (32, 4)),
+        ((32, 4), (16, 8)),
+        ((32, 4), (8, 16)),
+        ((32, 4), (4, 32)),
+        ((32, 4), (2, 64)),
+        ((16, 8), (64, 2)),
+        ((16, 8), (32, 4)),
+        ((16, 8), (16, 8)),
+        ((16, 8), (8, 16)),
+        ((16, 8), (4, 32)),
+        ((16, 8), (2, 64)),
+        ((8, 16), (64, 2)),
+        ((8, 16), (32, 4)),
+        ((8, 16), (16, 8)),
+        ((8, 16), (8, 16)),
+        ((8, 16), (4, 32)),
+        ((8, 16), (2, 64)),
+        ((4, 32), (64, 2)),
+        ((4, 32), (32, 4)),
+        ((4, 32), (16, 8)),
+        ((4, 32), (8, 16)),
+        ((4, 32), (4, 32)),
+        ((4, 32), (2, 64)),
+        ((2, 64), (64, 2)),
+        ((2, 64), (32, 4)),
+        ((2, 64), (16, 8)),
+        ((2, 64), (8, 16)),
+        ((2, 64), (4, 32)),
+        ((2, 64), (2, 64)),
+    ],
+    3: [
+        ((32, 2, 2), (32, 2, 2)),
+        ((32, 2, 2), (16, 4, 2)),
+        ((32, 2, 2), (16, 2, 4)),
+        ((32, 2, 2), (8, 8, 2)),
+        ((32, 2, 2), (8, 4, 4)),
+        ((32, 2, 2), (8, 2, 8)),
+        ((32, 2, 2), (4, 16, 2)),
+        ((32, 2, 2), (4, 8, 4)),
+        ((32, 2, 2), (4, 4, 8)),
+        ((32, 2, 2), (4, 2, 16)),
+        ((32, 2, 2), (2, 32, 2)),
+        ((32, 2, 2), (2, 16, 4)),
+        ((32, 2, 2), (2, 8, 8)),
+        ((32, 2, 2), (2, 4, 16)),
+        ((32, 2, 2), (2, 2, 32)),
+        ((16, 4, 2), (32, 2, 2)),
+        ((16, 4, 2), (16, 4, 2)),
+        ((16, 4, 2), (16, 2, 4)),
+        ((16, 4, 2), (8, 8, 2)),
+        ((16, 4, 2), (8, 4, 4)),
+        ((16, 4, 2), (8, 2, 8)),
+        ((16, 4, 2), (4, 16, 2)),
+        ((16, 4, 2), (4, 8, 4)),
+        ((16, 4, 2), (4, 4, 8)),
+        ((16, 4, 2), (4, 2, 16)),
+        ((16, 4, 2), (2, 32, 2)),
+        ((16, 4, 2), (2, 16, 4)),
+        ((16, 4, 2), (2, 8, 8)),
+        ((16, 4, 2), (2, 4, 16)),
+        ((16, 4, 2), (2, 2, 32)),
+        ((16, 2, 4), (32, 2, 2)),
+        ((16, 2, 4), (16, 4, 2)),
+        ((16, 2, 4), (16, 2, 4)),
+        ((16, 2, 4), (8, 8, 2)),
+        ((16, 2, 4), (8, 4, 4)),
+        ((16, 2, 4), (8, 2, 8)),
+        ((16, 2, 4), (4, 16, 2)),
+        ((16, 2, 4), (4, 8, 4)),
+        ((16, 2, 4), (4, 4, 8)),
+        ((16, 2, 4), (4, 2, 16)),
+        ((16, 2, 4), (2, 32, 2)),
+        ((16, 2, 4), (2, 16, 4)),
+        ((16, 2, 4), (2, 8, 8)),
+        ((16, 2, 4), (2, 4, 16)),
+        ((16, 2, 4), (2, 2, 32)),
+        ((8, 8, 2), (32, 2, 2)),
+        ((8, 8, 2), (16, 4, 2)),
+        ((8, 8, 2), (16, 2, 4)),
+        ((8, 8, 2), (8, 8, 2)),
+        ((8, 8, 2), (8, 4, 4)),
+        ((8, 8, 2), (8, 2, 8)),
+        ((8, 8, 2), (4, 16, 2)),
+        ((8, 8, 2), (4, 8, 4)),
+        ((8, 8, 2), (4, 4, 8)),
+        ((8, 8, 2), (4, 2, 16)),
+        ((8, 8, 2), (2, 32, 2)),
+        ((8, 8, 2), (2, 16, 4)),
+        ((8, 8, 2), (2, 8, 8)),
+        ((8, 8, 2), (2, 4, 16)),
+        ((8, 8, 2), (2, 2, 32)),
+        ((8, 4, 4), (32, 2, 2)),
+        ((8, 4, 4), (16, 4, 2)),
+        ((8, 4, 4), (16, 2, 4)),
+        ((8, 4, 4), (8, 8, 2)),
+        ((8, 4, 4), (8, 4, 4)),
+        ((8, 4, 4), (8, 2, 8)),
+        ((8, 4, 4), (4, 16, 2)),
+        ((8, 4, 4), (4, 8, 4)),
+        ((8, 4, 4), (4, 4, 8)),
+        ((8, 4, 4), (4, 2, 16)),
+        ((8, 4, 4), (2, 32, 2)),
+        ((8, 4, 4), (2, 16, 4)),
+        ((8, 4, 4), (2, 8, 8)),
+        ((8, 4, 4), (2, 4, 16)),
+        ((8, 4, 4), (2, 2, 32)),
+        ((8, 2, 8), (32, 2, 2)),
+        ((8, 2, 8), (16, 4, 2)),
+        ((8, 2, 8), (16, 2, 4)),
+        ((8, 2, 8), (8, 8, 2)),
+        ((8, 2, 8), (8, 4, 4)),
+        ((8, 2, 8), (8, 2, 8)),
+        ((8, 2, 8), (4, 16, 2)),
+        ((8, 2, 8), (4, 8, 4)),
+        ((8, 2, 8), (4, 4, 8)),
+        ((8, 2, 8), (4, 2, 16)),
+        ((8, 2, 8), (2, 32, 2)),
+        ((8, 2, 8), (2, 16, 4)),
+        ((8, 2, 8), (2, 8, 8)),
+        ((8, 2, 8), (2, 4, 16)),
+        ((8, 2, 8), (2, 2, 32)),
+        ((4, 16, 2), (32, 2, 2)),
+        ((4, 16, 2), (16, 4, 2)),
+        ((4, 16, 2), (16, 2, 4)),
+        ((4, 16, 2), (8, 8, 2)),
+        ((4, 16, 2), (8, 4, 4)),
+        ((4, 16, 2), (8, 2, 8)),
+        ((4, 16, 2), (4, 16, 2)),
+        ((4, 16, 2), (4, 8, 4)),
+        ((4, 16, 2), (4, 4, 8)),
+        ((4, 16, 2), (4, 2, 16)),
+        ((4, 16, 2), (2, 32, 2)),
+        ((4, 16, 2), (2, 16, 4)),
+        ((4, 16, 2), (2, 8, 8)),
+        ((4, 16, 2), (2, 4, 16)),
+        ((4, 16, 2), (2, 2, 32)),
+        ((4, 8, 4), (32, 2, 2)),
+        ((4, 8, 4), (16, 4, 2)),
+        ((4, 8, 4), (16, 2, 4)),
+        ((4, 8, 4), (8, 8, 2)),
+        ((4, 8, 4), (8, 4, 4)),
+        ((4, 8, 4), (8, 2, 8)),
+        ((4, 8, 4), (4, 16, 2)),
+        ((4, 8, 4), (4, 8, 4)),
+        ((4, 8, 4), (4, 4, 8)),
+        ((4, 8, 4), (4, 2, 16)),
+        ((4, 8, 4), (2, 32, 2)),
+        ((4, 8, 4), (2, 16, 4)),
+        ((4, 8, 4), (2, 8, 8)),
+        ((4, 8, 4), (2, 4, 16)),
+        ((4, 8, 4), (2, 2, 32)),
+        ((4, 4, 8), (32, 2, 2)),
+        ((4, 4, 8), (16, 4, 2)),
+        ((4, 4, 8), (16, 2, 4)),
+        ((4, 4, 8), (8, 8, 2)),
+        ((4, 4, 8), (8, 4, 4)),
+        ((4, 4, 8), (8, 2, 8)),
+        ((4, 4, 8), (4, 16, 2)),
+        ((4, 4, 8), (4, 8, 4)),
+        ((4, 4, 8), (4, 4, 8)),
+        ((4, 4, 8), (4, 2, 16)),
+        ((4, 4, 8), (2, 32, 2)),
+        ((4, 4, 8), (2, 16, 4)),
+        ((4, 4, 8), (2, 8, 8)),
+        ((4, 4, 8), (2, 4, 16)),
+        ((4, 4, 8), (2, 2, 32)),
+        ((4, 2, 16), (32, 2, 2)),
+        ((4, 2, 16), (16, 4, 2)),
+        ((4, 2, 16), (16, 2, 4)),
+        ((4, 2, 16), (8, 8, 2)),
+        ((4, 2, 16), (8, 4, 4)),
+        ((4, 2, 16), (8, 2, 8)),
+        ((4, 2, 16), (4, 16, 2)),
+        ((4, 2, 16), (4, 8, 4)),
+        ((4, 2, 16), (4, 4, 8)),
+        ((4, 2, 16), (4, 2, 16)),
+        ((4, 2, 16), (2, 32, 2)),
+        ((4, 2, 16), (2, 16, 4)),
+        ((4, 2, 16), (2, 8, 8)),
+        ((4, 2, 16), (2, 4, 16)),
+        ((4, 2, 16), (2, 2, 32)),
+        ((2, 32, 2), (32, 2, 2)),
+        ((2, 32, 2), (16, 4, 2)),
+        ((2, 32, 2), (16, 2, 4)),
+        ((2, 32, 2), (8, 8, 2)),
+        ((2, 32, 2), (8, 4, 4)),
+        ((2, 32, 2), (8, 2, 8)),
+        ((2, 32, 2), (4, 16, 2)),
+        ((2, 32, 2), (4, 8, 4)),
+        ((2, 32, 2), (4, 4, 8)),
+        ((2, 32, 2), (4, 2, 16)),
+        ((2, 32, 2), (2, 32, 2)),
+        ((2, 32, 2), (2, 16, 4)),
+        ((2, 32, 2), (2, 8, 8)),
+        ((2, 32, 2), (2, 4, 16)),
+        ((2, 32, 2), (2, 2, 32)),
+        ((2, 16, 4), (32, 2, 2)),
+        ((2, 16, 4), (16, 4, 2)),
+        ((2, 16, 4), (16, 2, 4)),
+        ((2, 16, 4), (8, 8, 2)),
+        ((2, 16, 4), (8, 4, 4)),
+        ((2, 16, 4), (8, 2, 8)),
+        ((2, 16, 4), (4, 16, 2)),
+        ((2, 16, 4), (4, 8, 4)),
+        ((2, 16, 4), (4, 4, 8)),
+        ((2, 16, 4), (4, 2, 16)),
+        ((2, 16, 4), (2, 32, 2)),
+        ((2, 16, 4), (2, 16, 4)),
+        ((2, 16, 4), (2, 8, 8)),
+        ((2, 16, 4), (2, 4, 16)),
+        ((2, 16, 4), (2, 2, 32)),
+        ((2, 8, 8), (32, 2, 2)),
+        ((2, 8, 8), (16, 4, 2)),
+        ((2, 8, 8), (16, 2, 4)),
+        ((2, 8, 8), (8, 8, 2)),
+        ((2, 8, 8), (8, 4, 4)),
+        ((2, 8, 8), (8, 2, 8)),
+        ((2, 8, 8), (4, 16, 2)),
+        ((2, 8, 8), (4, 8, 4)),
+        ((2, 8, 8), (4, 4, 8)),
+        ((2, 8, 8), (4, 2, 16)),
+        ((2, 8, 8), (2, 32, 2)),
+        ((2, 8, 8), (2, 16, 4)),
+        ((2, 8, 8), (2, 8, 8)),
+        ((2, 8, 8), (2, 4, 16)),
+        ((2, 8, 8), (2, 2, 32)),
+        ((2, 4, 16), (32, 2, 2)),
+        ((2, 4, 16), (16, 4, 2)),
+        ((2, 4, 16), (16, 2, 4)),
+        ((2, 4, 16), (8, 8, 2)),
+        ((2, 4, 16), (8, 4, 4)),
+        ((2, 4, 16), (8, 2, 8)),
+        ((2, 4, 16), (4, 16, 2)),
+        ((2, 4, 16), (4, 8, 4)),
+        ((2, 4, 16), (4, 4, 8)),
+        ((2, 4, 16), (4, 2, 16)),
+        ((2, 4, 16), (2, 32, 2)),
+        ((2, 4, 16), (2, 16, 4)),
+        ((2, 4, 16), (2, 8, 8)),
+        ((2, 4, 16), (2, 4, 16)),
+        ((2, 4, 16), (2, 2, 32)),
+        ((2, 2, 32), (32, 2, 2)),
+        ((2, 2, 32), (16, 4, 2)),
+        ((2, 2, 32), (16, 2, 4)),
+        ((2, 2, 32), (8, 8, 2)),
+        ((2, 2, 32), (8, 4, 4)),
+        ((2, 2, 32), (8, 2, 8)),
+        ((2, 2, 32), (4, 16, 2)),
+        ((2, 2, 32), (4, 8, 4)),
+        ((2, 2, 32), (4, 4, 8)),
+        ((2, 2, 32), (4, 2, 16)),
+        ((2, 2, 32), (2, 32, 2)),
+        ((2, 2, 32), (2, 16, 4)),
+        ((2, 2, 32), (2, 8, 8)),
+        ((2, 2, 32), (2, 4, 16)),
+        ((2, 2, 32), (2, 2, 32)),
+    ],
+}

build/torch210-cu126-aarch64-linux/backends/configs/cutlass/fna_backward_128x64.py ADDED Viewed

	@@ -0,0 +1,223 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+from typing import Dict, List
+from ...._types import CutlassFnaForwardConfigType
+# NOTE: we're excluding tile shapes that include 1 just to
+# reduce the giant number of configs down to a reasonable
+# amount; otherwise autotuning would take more than a few
+# seconds per call which is unacceptable. Tile shapes with
+# 1s are rarely selected.
+_FNA_BACKWARD_128x64_TILE_SIZES: Dict[int, List[CutlassFnaForwardConfigType]] = {
+    1: [
+        ((128,), (64,)),
+    ],
+    2: [
+        ((64, 2), (32, 2)),
+        ((64, 2), (16, 4)),
+        ((64, 2), (8, 8)),
+        ((64, 2), (4, 16)),
+        ((64, 2), (2, 32)),
+        ((32, 4), (32, 2)),
+        ((32, 4), (16, 4)),
+        ((32, 4), (8, 8)),
+        ((32, 4), (4, 16)),
+        ((32, 4), (2, 32)),
+        ((16, 8), (32, 2)),
+        ((16, 8), (16, 4)),
+        ((16, 8), (8, 8)),
+        ((16, 8), (4, 16)),
+        ((16, 8), (2, 32)),
+        ((8, 16), (32, 2)),
+        ((8, 16), (16, 4)),
+        ((8, 16), (8, 8)),
+        ((8, 16), (4, 16)),
+        ((8, 16), (2, 32)),
+        ((4, 32), (32, 2)),
+        ((4, 32), (16, 4)),
+        ((4, 32), (8, 8)),
+        ((4, 32), (4, 16)),
+        ((4, 32), (2, 32)),
+        ((2, 64), (32, 2)),
+        ((2, 64), (16, 4)),
+        ((2, 64), (8, 8)),
+        ((2, 64), (4, 16)),
+        ((2, 64), (2, 32)),
+    ],
+    3: [
+        ((32, 2, 2), (16, 2, 2)),
+        ((32, 2, 2), (8, 4, 2)),
+        ((32, 2, 2), (8, 2, 4)),
+        ((32, 2, 2), (4, 8, 2)),
+        ((32, 2, 2), (4, 4, 4)),
+        ((32, 2, 2), (4, 2, 8)),
+        ((32, 2, 2), (2, 16, 2)),
+        ((32, 2, 2), (2, 8, 4)),
+        ((32, 2, 2), (2, 4, 8)),
+        ((32, 2, 2), (2, 2, 16)),
+        ((16, 4, 2), (16, 2, 2)),
+        ((16, 4, 2), (8, 4, 2)),
+        ((16, 4, 2), (8, 2, 4)),
+        ((16, 4, 2), (4, 8, 2)),
+        ((16, 4, 2), (4, 4, 4)),
+        ((16, 4, 2), (4, 2, 8)),
+        ((16, 4, 2), (2, 16, 2)),
+        ((16, 4, 2), (2, 8, 4)),
+        ((16, 4, 2), (2, 4, 8)),
+        ((16, 4, 2), (2, 2, 16)),
+        ((16, 2, 4), (16, 2, 2)),
+        ((16, 2, 4), (8, 4, 2)),
+        ((16, 2, 4), (8, 2, 4)),
+        ((16, 2, 4), (4, 8, 2)),
+        ((16, 2, 4), (4, 4, 4)),
+        ((16, 2, 4), (4, 2, 8)),
+        ((16, 2, 4), (2, 16, 2)),
+        ((16, 2, 4), (2, 8, 4)),
+        ((16, 2, 4), (2, 4, 8)),
+        ((16, 2, 4), (2, 2, 16)),
+        ((8, 8, 2), (16, 2, 2)),
+        ((8, 8, 2), (8, 4, 2)),
+        ((8, 8, 2), (8, 2, 4)),
+        ((8, 8, 2), (4, 8, 2)),
+        ((8, 8, 2), (4, 4, 4)),
+        ((8, 8, 2), (4, 2, 8)),
+        ((8, 8, 2), (2, 16, 2)),
+        ((8, 8, 2), (2, 8, 4)),
+        ((8, 8, 2), (2, 4, 8)),
+        ((8, 8, 2), (2, 2, 16)),
+        ((8, 4, 4), (16, 2, 2)),
+        ((8, 4, 4), (8, 4, 2)),
+        ((8, 4, 4), (8, 2, 4)),
+        ((8, 4, 4), (4, 8, 2)),
+        ((8, 4, 4), (4, 4, 4)),
+        ((8, 4, 4), (4, 2, 8)),
+        ((8, 4, 4), (2, 16, 2)),
+        ((8, 4, 4), (2, 8, 4)),
+        ((8, 4, 4), (2, 4, 8)),
+        ((8, 4, 4), (2, 2, 16)),
+        ((8, 2, 8), (16, 2, 2)),
+        ((8, 2, 8), (8, 4, 2)),
+        ((8, 2, 8), (8, 2, 4)),
+        ((8, 2, 8), (4, 8, 2)),
+        ((8, 2, 8), (4, 4, 4)),
+        ((8, 2, 8), (4, 2, 8)),
+        ((8, 2, 8), (2, 16, 2)),
+        ((8, 2, 8), (2, 8, 4)),
+        ((8, 2, 8), (2, 4, 8)),
+        ((8, 2, 8), (2, 2, 16)),
+        ((4, 16, 2), (16, 2, 2)),
+        ((4, 16, 2), (8, 4, 2)),
+        ((4, 16, 2), (8, 2, 4)),
+        ((4, 16, 2), (4, 8, 2)),
+        ((4, 16, 2), (4, 4, 4)),
+        ((4, 16, 2), (4, 2, 8)),
+        ((4, 16, 2), (2, 16, 2)),
+        ((4, 16, 2), (2, 8, 4)),
+        ((4, 16, 2), (2, 4, 8)),
+        ((4, 16, 2), (2, 2, 16)),
+        ((4, 8, 4), (16, 2, 2)),
+        ((4, 8, 4), (8, 4, 2)),
+        ((4, 8, 4), (8, 2, 4)),
+        ((4, 8, 4), (4, 8, 2)),
+        ((4, 8, 4), (4, 4, 4)),
+        ((4, 8, 4), (4, 2, 8)),
+        ((4, 8, 4), (2, 16, 2)),
+        ((4, 8, 4), (2, 8, 4)),
+        ((4, 8, 4), (2, 4, 8)),
+        ((4, 8, 4), (2, 2, 16)),
+        ((4, 4, 8), (16, 2, 2)),
+        ((4, 4, 8), (8, 4, 2)),
+        ((4, 4, 8), (8, 2, 4)),
+        ((4, 4, 8), (4, 8, 2)),
+        ((4, 4, 8), (4, 4, 4)),
+        ((4, 4, 8), (4, 2, 8)),
+        ((4, 4, 8), (2, 16, 2)),
+        ((4, 4, 8), (2, 8, 4)),
+        ((4, 4, 8), (2, 4, 8)),
+        ((4, 4, 8), (2, 2, 16)),
+        ((4, 2, 16), (16, 2, 2)),
+        ((4, 2, 16), (8, 4, 2)),
+        ((4, 2, 16), (8, 2, 4)),
+        ((4, 2, 16), (4, 8, 2)),
+        ((4, 2, 16), (4, 4, 4)),
+        ((4, 2, 16), (4, 2, 8)),
+        ((4, 2, 16), (2, 16, 2)),
+        ((4, 2, 16), (2, 8, 4)),
+        ((4, 2, 16), (2, 4, 8)),
+        ((4, 2, 16), (2, 2, 16)),
+        ((2, 32, 2), (16, 2, 2)),
+        ((2, 32, 2), (8, 4, 2)),
+        ((2, 32, 2), (8, 2, 4)),
+        ((2, 32, 2), (4, 8, 2)),
+        ((2, 32, 2), (4, 4, 4)),
+        ((2, 32, 2), (4, 2, 8)),
+        ((2, 32, 2), (2, 16, 2)),
+        ((2, 32, 2), (2, 8, 4)),
+        ((2, 32, 2), (2, 4, 8)),
+        ((2, 32, 2), (2, 2, 16)),
+        ((2, 16, 4), (16, 2, 2)),
+        ((2, 16, 4), (8, 4, 2)),
+        ((2, 16, 4), (8, 2, 4)),
+        ((2, 16, 4), (4, 8, 2)),
+        ((2, 16, 4), (4, 4, 4)),
+        ((2, 16, 4), (4, 2, 8)),
+        ((2, 16, 4), (2, 16, 2)),
+        ((2, 16, 4), (2, 8, 4)),
+        ((2, 16, 4), (2, 4, 8)),
+        ((2, 16, 4), (2, 2, 16)),
+        ((2, 8, 8), (16, 2, 2)),
+        ((2, 8, 8), (8, 4, 2)),
+        ((2, 8, 8), (8, 2, 4)),
+        ((2, 8, 8), (4, 8, 2)),
+        ((2, 8, 8), (4, 4, 4)),
+        ((2, 8, 8), (4, 2, 8)),
+        ((2, 8, 8), (2, 16, 2)),
+        ((2, 8, 8), (2, 8, 4)),
+        ((2, 8, 8), (2, 4, 8)),
+        ((2, 8, 8), (2, 2, 16)),
+        ((2, 4, 16), (16, 2, 2)),
+        ((2, 4, 16), (8, 4, 2)),
+        ((2, 4, 16), (8, 2, 4)),
+        ((2, 4, 16), (4, 8, 2)),
+        ((2, 4, 16), (4, 4, 4)),
+        ((2, 4, 16), (4, 2, 8)),
+        ((2, 4, 16), (2, 16, 2)),
+        ((2, 4, 16), (2, 8, 4)),
+        ((2, 4, 16), (2, 4, 8)),
+        ((2, 4, 16), (2, 2, 16)),
+        ((2, 2, 32), (16, 2, 2)),
+        ((2, 2, 32), (8, 4, 2)),
+        ((2, 2, 32), (8, 2, 4)),
+        ((2, 2, 32), (4, 8, 2)),
+        ((2, 2, 32), (4, 4, 4)),
+        ((2, 2, 32), (4, 2, 8)),
+        ((2, 2, 32), (2, 16, 2)),
+        ((2, 2, 32), (2, 8, 4)),
+        ((2, 2, 32), (2, 4, 8)),
+        ((2, 2, 32), (2, 2, 16)),
+    ],
+}

build/torch210-cu126-aarch64-linux/backends/configs/cutlass/fna_backward_64x64.py ADDED Viewed

	@@ -0,0 +1,168 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+from typing import Dict, List
+from ...._types import QKTileShapeType
+# NOTE: we're excluding tile shapes that include 1 just to
+# reduce the giant number of configs down to a reasonable
+# amount; otherwise autotuning would take more than a few
+# seconds per call which is unacceptable. Tile shapes with
+# 1s are rarely selected.
+_FNA_BACKWARD_64x64_TILE_SIZES: Dict[int, List[QKTileShapeType]] = {
+    1: [
+        ((64,), (64,)),
+    ],
+    2: [
+        ((32, 2), (32, 2)),
+        ((32, 2), (16, 4)),
+        ((32, 2), (8, 8)),
+        ((32, 2), (4, 16)),
+        ((32, 2), (2, 32)),
+        ((16, 4), (32, 2)),
+        ((16, 4), (16, 4)),
+        ((16, 4), (8, 8)),
+        ((16, 4), (4, 16)),
+        ((16, 4), (2, 32)),
+        ((8, 8), (32, 2)),
+        ((8, 8), (16, 4)),
+        ((8, 8), (8, 8)),
+        ((8, 8), (4, 16)),
+        ((8, 8), (2, 32)),
+        ((4, 16), (32, 2)),
+        ((4, 16), (16, 4)),
+        ((4, 16), (8, 8)),
+        ((4, 16), (4, 16)),
+        ((4, 16), (2, 32)),
+        ((2, 32), (32, 2)),
+        ((2, 32), (16, 4)),
+        ((2, 32), (8, 8)),
+        ((2, 32), (4, 16)),
+        ((2, 32), (2, 32)),
+    ],
+    3: [
+        ((16, 2, 2), (16, 2, 2)),
+        ((16, 2, 2), (8, 4, 2)),
+        ((16, 2, 2), (8, 2, 4)),
+        ((16, 2, 2), (4, 8, 2)),
+        ((16, 2, 2), (4, 4, 4)),
+        ((16, 2, 2), (4, 2, 8)),
+        ((16, 2, 2), (2, 16, 2)),
+        ((16, 2, 2), (2, 8, 4)),
+        ((16, 2, 2), (2, 4, 8)),
+        ((16, 2, 2), (2, 2, 16)),
+        ((8, 4, 2), (16, 2, 2)),
+        ((8, 4, 2), (8, 4, 2)),
+        ((8, 4, 2), (8, 2, 4)),
+        ((8, 4, 2), (4, 8, 2)),
+        ((8, 4, 2), (4, 4, 4)),
+        ((8, 4, 2), (4, 2, 8)),
+        ((8, 4, 2), (2, 16, 2)),
+        ((8, 4, 2), (2, 8, 4)),
+        ((8, 4, 2), (2, 4, 8)),
+        ((8, 4, 2), (2, 2, 16)),
+        ((8, 2, 4), (16, 2, 2)),
+        ((8, 2, 4), (8, 4, 2)),
+        ((8, 2, 4), (8, 2, 4)),
+        ((8, 2, 4), (4, 8, 2)),
+        ((8, 2, 4), (4, 4, 4)),
+        ((8, 2, 4), (4, 2, 8)),
+        ((8, 2, 4), (2, 16, 2)),
+        ((8, 2, 4), (2, 8, 4)),
+        ((8, 2, 4), (2, 4, 8)),
+        ((8, 2, 4), (2, 2, 16)),
+        ((4, 8, 2), (16, 2, 2)),
+        ((4, 8, 2), (8, 4, 2)),
+        ((4, 8, 2), (8, 2, 4)),
+        ((4, 8, 2), (4, 8, 2)),
+        ((4, 8, 2), (4, 4, 4)),
+        ((4, 8, 2), (4, 2, 8)),
+        ((4, 8, 2), (2, 16, 2)),
+        ((4, 8, 2), (2, 8, 4)),
+        ((4, 8, 2), (2, 4, 8)),
+        ((4, 8, 2), (2, 2, 16)),
+        ((4, 4, 4), (16, 2, 2)),
+        ((4, 4, 4), (8, 4, 2)),
+        ((4, 4, 4), (8, 2, 4)),
+        ((4, 4, 4), (4, 8, 2)),
+        ((4, 4, 4), (4, 4, 4)),
+        ((4, 4, 4), (4, 2, 8)),
+        ((4, 4, 4), (2, 16, 2)),
+        ((4, 4, 4), (2, 8, 4)),
+        ((4, 4, 4), (2, 4, 8)),
+        ((4, 4, 4), (2, 2, 16)),
+        ((4, 2, 8), (16, 2, 2)),
+        ((4, 2, 8), (8, 4, 2)),
+        ((4, 2, 8), (8, 2, 4)),
+        ((4, 2, 8), (4, 8, 2)),
+        ((4, 2, 8), (4, 4, 4)),
+        ((4, 2, 8), (4, 2, 8)),
+        ((4, 2, 8), (2, 16, 2)),
+        ((4, 2, 8), (2, 8, 4)),
+        ((4, 2, 8), (2, 4, 8)),
+        ((4, 2, 8), (2, 2, 16)),
+        ((2, 16, 2), (16, 2, 2)),
+        ((2, 16, 2), (8, 4, 2)),
+        ((2, 16, 2), (8, 2, 4)),
+        ((2, 16, 2), (4, 8, 2)),
+        ((2, 16, 2), (4, 4, 4)),
+        ((2, 16, 2), (4, 2, 8)),
+        ((2, 16, 2), (2, 16, 2)),
+        ((2, 16, 2), (2, 8, 4)),
+        ((2, 16, 2), (2, 4, 8)),
+        ((2, 16, 2), (2, 2, 16)),
+        ((2, 8, 4), (16, 2, 2)),
+        ((2, 8, 4), (8, 4, 2)),
+        ((2, 8, 4), (8, 2, 4)),
+        ((2, 8, 4), (4, 8, 2)),
+        ((2, 8, 4), (4, 4, 4)),
+        ((2, 8, 4), (4, 2, 8)),
+        ((2, 8, 4), (2, 16, 2)),
+        ((2, 8, 4), (2, 8, 4)),
+        ((2, 8, 4), (2, 4, 8)),
+        ((2, 8, 4), (2, 2, 16)),
+        ((2, 4, 8), (16, 2, 2)),
+        ((2, 4, 8), (8, 4, 2)),
+        ((2, 4, 8), (8, 2, 4)),
+        ((2, 4, 8), (4, 8, 2)),
+        ((2, 4, 8), (4, 4, 4)),
+        ((2, 4, 8), (4, 2, 8)),
+        ((2, 4, 8), (2, 16, 2)),
+        ((2, 4, 8), (2, 8, 4)),
+        ((2, 4, 8), (2, 4, 8)),
+        ((2, 4, 8), (2, 2, 16)),
+        ((2, 2, 16), (16, 2, 2)),
+        ((2, 2, 16), (8, 4, 2)),
+        ((2, 2, 16), (8, 2, 4)),
+        ((2, 2, 16), (4, 8, 2)),
+        ((2, 2, 16), (4, 4, 4)),
+        ((2, 2, 16), (4, 2, 8)),
+        ((2, 2, 16), (2, 16, 2)),
+        ((2, 2, 16), (2, 8, 4)),
+        ((2, 2, 16), (2, 4, 8)),
+        ((2, 2, 16), (2, 2, 16)),
+    ],
+}

build/torch210-cu126-aarch64-linux/backends/configs/cutlass/fna_forward_32x128.py ADDED Viewed

	@@ -0,0 +1,90 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+from typing import Dict, List
+from ...._types import QKTileShapeType
+# TODO: More combinations are possible for
+# 2D and 3D (query tile does not have to be smaller
+# than KV tile); but that behavior is untested,
+# and IIRC was unstable.
+_FNA_FORWARD_32x128_TILE_SIZES: Dict[int, List[QKTileShapeType]] = {
+    1: [
+        ((32,), (128,)),
+    ],
+    2: [
+        ((16, 2), (64, 2)),
+        ((16, 2), (32, 4)),
+        ((16, 2), (16, 8)),
+        ((8, 4), (32, 4)),
+        ((8, 4), (16, 8)),
+        ((8, 4), (8, 16)),
+        ((4, 8), (16, 8)),
+        ((4, 8), (8, 16)),
+        ((4, 8), (4, 32)),
+        ((2, 16), (8, 16)),
+        ((2, 16), (4, 32)),
+        ((2, 16), (2, 64)),
+    ],
+    3: [
+        ((8, 2, 2), (32, 2, 2)),
+        ((8, 2, 2), (16, 4, 2)),
+        ((8, 2, 2), (16, 2, 4)),
+        ((8, 2, 2), (8, 8, 2)),
+        ((8, 2, 2), (8, 4, 4)),
+        ((8, 2, 2), (8, 2, 8)),
+        ((4, 4, 2), (16, 4, 2)),
+        ((4, 4, 2), (8, 8, 2)),
+        ((4, 4, 2), (8, 4, 4)),
+        ((4, 4, 2), (4, 16, 2)),
+        ((4, 4, 2), (4, 8, 4)),
+        ((4, 4, 2), (4, 4, 8)),
+        ((4, 2, 4), (16, 2, 4)),
+        ((4, 2, 4), (8, 4, 4)),
+        ((4, 2, 4), (8, 2, 8)),
+        ((4, 2, 4), (4, 8, 4)),
+        ((4, 2, 4), (4, 4, 8)),
+        ((4, 2, 4), (4, 2, 16)),
+        ((2, 8, 2), (8, 8, 2)),
+        ((2, 8, 2), (4, 16, 2)),
+        ((2, 8, 2), (4, 8, 4)),
+        ((2, 8, 2), (2, 32, 2)),
+        ((2, 8, 2), (2, 16, 4)),
+        ((2, 8, 2), (2, 8, 8)),
+        ((2, 4, 4), (8, 4, 4)),
+        ((2, 4, 4), (4, 8, 4)),
+        ((2, 4, 4), (4, 4, 8)),
+        ((2, 4, 4), (2, 16, 4)),
+        ((2, 4, 4), (2, 8, 8)),
+        ((2, 4, 4), (2, 4, 16)),
+        ((2, 2, 8), (8, 2, 8)),
+        ((2, 2, 8), (4, 4, 8)),
+        ((2, 2, 8), (4, 2, 16)),
+        ((2, 2, 8), (2, 8, 8)),
+        ((2, 2, 8), (2, 4, 16)),
+        ((2, 2, 8), (2, 2, 32)),
+    ],
+}

build/torch210-cu126-aarch64-linux/backends/configs/cutlass/fna_forward_64x128.py ADDED Viewed

	@@ -0,0 +1,82 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+from typing import Dict, List
+from ...._types import QKTileShapeType
+# TODO: More combinations are possible for
+# 2D and 3D (query tile does not have to be smaller
+# than KV tile); but that behavior is untested,
+# and IIRC was unstable.
+_FNA_FORWARD_64x128_TILE_SIZES: Dict[int, List[QKTileShapeType]] = {
+    1: [
+        ((64,), (128,)),
+    ],
+    2: [
+        ((32, 2), (64, 2)),
+        ((32, 2), (32, 4)),
+        ((16, 4), (32, 4)),
+        ((16, 4), (16, 8)),
+        ((8, 8), (16, 8)),
+        ((8, 8), (8, 16)),
+        ((4, 16), (8, 16)),
+        ((4, 16), (4, 32)),
+        ((2, 32), (4, 32)),
+        ((2, 32), (2, 64)),
+    ],
+    3: [
+        ((16, 2, 2), (32, 2, 2)),
+        ((16, 2, 2), (16, 4, 2)),
+        ((16, 2, 2), (16, 2, 4)),
+        ((8, 4, 2), (16, 4, 2)),
+        ((8, 4, 2), (8, 8, 2)),
+        ((8, 4, 2), (8, 4, 4)),
+        ((8, 2, 4), (16, 2, 4)),
+        ((8, 2, 4), (8, 4, 4)),
+        ((8, 2, 4), (8, 2, 8)),
+        ((4, 8, 2), (8, 8, 2)),
+        ((4, 8, 2), (4, 16, 2)),
+        ((4, 8, 2), (4, 8, 4)),
+        ((4, 4, 4), (8, 4, 4)),
+        ((4, 4, 4), (4, 8, 4)),
+        ((4, 4, 4), (4, 4, 8)),
+        ((4, 2, 8), (8, 2, 8)),
+        ((4, 2, 8), (4, 4, 8)),
+        ((4, 2, 8), (4, 2, 16)),
+        ((2, 16, 2), (4, 16, 2)),
+        ((2, 16, 2), (2, 32, 2)),
+        ((2, 16, 2), (2, 16, 4)),
+        ((2, 8, 4), (4, 8, 4)),
+        ((2, 8, 4), (2, 16, 4)),
+        ((2, 8, 4), (2, 8, 8)),
+        ((2, 4, 8), (4, 4, 8)),
+        ((2, 4, 8), (2, 8, 8)),
+        ((2, 4, 8), (2, 4, 16)),
+        ((2, 2, 16), (4, 2, 16)),
+        ((2, 2, 16), (2, 4, 16)),
+        ((2, 2, 16), (2, 2, 32)),
+    ],
+}

build/torch210-cu126-aarch64-linux/backends/configs/cutlass/fna_forward_64x64.py ADDED Viewed

	@@ -0,0 +1,63 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+from typing import Dict, List
+from ...._types import QKTileShapeType
+# TODO: More combinations are possible for
+# 2D and 3D (query tile does not have to be smaller
+# than KV tile); but that behavior is untested,
+# and IIRC was unstable.
+# NOTE: we're excluding tile shapes that include 1 just to
+# reduce the giant number of configs down to a reasonable
+# amount; otherwise autotuning would take more than a few
+# seconds per call which is unacceptable. Tile shapes with
+# 1s are rarely selected.
+_FNA_FORWARD_64x64_TILE_SIZES: Dict[int, List[QKTileShapeType]] = {
+    1: [
+        ((64,), (64,)),
+    ],
+    2: [
+        ((32, 2), (32, 2)),
+        ((16, 4), (16, 4)),
+        ((8, 8), (8, 8)),
+        ((4, 16), (4, 16)),
+        ((2, 32), (2, 32)),
+    ],
+    3: [
+        ((16, 2, 2), (16, 2, 2)),
+        ((8, 4, 2), (8, 4, 2)),
+        ((8, 2, 4), (8, 2, 4)),
+        ((4, 8, 2), (4, 8, 2)),
+        ((4, 4, 4), (4, 4, 4)),
+        ((4, 2, 8), (4, 2, 8)),
+        ((2, 16, 2), (2, 16, 2)),
+        ((2, 8, 4), (2, 8, 4)),
+        ((2, 4, 8), (2, 4, 8)),
+        ((2, 2, 16), (2, 2, 16)),
+    ],
+}

build/torch210-cu126-aarch64-linux/backends/configs/cutlass_blackwell/__init__.py ADDED Viewed

	@@ -0,0 +1,391 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+from typing import List, Optional
+import torch  # noqa: F401
+from torch import Tensor
+from ...._types import (
+    CutlassBlackwellFmhaBackwardConfigType,
+    CutlassBlackwellFmhaForwardConfigType,
+    CutlassBlackwellFnaBackwardConfigType,
+    CutlassBlackwellFnaForwardConfigType,
+    DimensionType,
+)
+from ....utils.checks import check_tile_shape
+from ....utils.device import get_device_cc
+# The current CUTLASS FMHA forward kernel can only do Q tile size 256, KV tile size 128.
+# This limits 1D tile shapes to just the one, but for 2-D and 3-D we can have many more shapes,
+# only some of which we compile. Adding new ones requires adding them to autogen, regenerating
+# the instantiations, and recompiling libnatten. Unlike CUTLASS 2.X FNA, multi-dim tile shapes are
+# static in Blackwell FNA, and not dynamic.
+BLACKWELL_FORWARD_TILE_SHAPES = {
+    1: [
+        ((256,), (128,)),
+    ],
+    2: [
+        ((16, 16), (16, 8)),
+        ((16, 16), (8, 16)),
+        ((8, 32), (8, 16)),
+        ((8, 32), (4, 32)),
+    ],
+    3: [
+        ((8, 4, 8), (4, 4, 8)),
+        ((8, 4, 8), (2, 8, 8)),
+        ((2, 8, 16), (4, 4, 8)),
+        ((2, 8, 16), (2, 8, 8)),
+        ((4, 4, 16), (2, 4, 16)),
+        ((2, 16, 8), (2, 8, 8)),
+        ((4, 8, 8), (2, 8, 8)),
+    ],
+}
+BLACKWELL_BACKWARD_TILE_SHAPES = {
+    1: [
+        ((128,), (128,)),
+    ],
+    2: [
+        ((16, 8), (16, 8)),
+        ((16, 8), (8, 16)),
+        ((8, 16), (16, 8)),
+        ((8, 16), (8, 16)),
+    ],
+    3: [
+        ((4, 4, 8), (4, 4, 8)),
+        ((4, 4, 8), (2, 8, 8)),
+        ((1, 8, 16), (4, 4, 8)),
+        ((2, 8, 8), (4, 4, 8)),
+        ((1, 8, 16), (2, 8, 8)),
+        ((2, 4, 16), (2, 4, 16)),
+        ((4, 2, 16), (2, 4, 16)),
+        ((4, 4, 8), (2, 4, 16)),
+        ((2, 8, 8), (2, 8, 8)),
+    ],
+}
+def _get_default_tile_shapes_forward(
+    na_dim: int,
+) -> CutlassBlackwellFnaForwardConfigType:
+    assert na_dim in [1, 2, 3]
+    if na_dim == 1:
+        return ((256,), (128,))
+    if na_dim == 2:
+        return ((16, 16), (16, 8))
+    if na_dim == 3:
+        return ((8, 4, 8), (4, 4, 8))
+    raise NotImplementedError()
+def get_all_forward_configs(
+    input_tensor: Tensor,
+) -> List[CutlassBlackwellFnaForwardConfigType]:
+    assert input_tensor.dim() in [4, 5, 6]
+    na_dim = input_tensor.dim() - 3  # batch, heads, head_dim
+    device_cc = get_device_cc(input_tensor.device)
+    if device_cc not in [100, 103]:
+        return []
+    return BLACKWELL_FORWARD_TILE_SHAPES[na_dim]  # type: ignore
+def get_all_backward_configs(
+    input_tensor: Tensor,
+) -> List[CutlassBlackwellFnaBackwardConfigType]:
+    assert input_tensor.dim() in [4, 5, 6]
+    na_dim = input_tensor.dim() - 3  # batch, heads, head_dim
+    device_cc = get_device_cc(input_tensor.device)
+    if device_cc not in [100, 103]:
+        return []
+    return BLACKWELL_BACKWARD_TILE_SHAPES[na_dim]  # type: ignore
+# For FMHA
+def get_all_fmha_forward_configs(
+    input_tensor: Tensor,
+) -> List[CutlassBlackwellFmhaForwardConfigType]:
+    if input_tensor.dim() != 4:
+        raise ValueError("Only 4-D tensors are supported in FMHA.")
+    tile_shapes = get_all_forward_configs(input_tensor)
+    assert all(len(q_t) == len(kv_t) == 1 for q_t, kv_t in tile_shapes)
+    tile_sizes = [(q_t[0], kv_t[0]) for q_t, kv_t in tile_shapes]
+    return tile_sizes
+def get_all_fmha_backward_configs(
+    input_tensor: Tensor,
+) -> List[CutlassBlackwellFmhaBackwardConfigType]:
+    if input_tensor.dim() != 4:
+        raise ValueError("Only 4-D tensors are supported in FMHA.")
+    tile_shapes = get_all_backward_configs(input_tensor)
+    assert all(len(q_t) == len(kv_t) == 1 for q_t, kv_t in tile_shapes)
+    tile_sizes = [(q_t[0], kv_t[0]) for q_t, kv_t in tile_shapes]
+    return tile_sizes
+def get_default_forward_tile_shapes(
+    input_tensor: Tensor,
+) -> CutlassBlackwellFnaForwardConfigType:
+    assert input_tensor.dim() in [4, 5, 6]
+    na_dim = input_tensor.dim() - 3  # batch, heads, head_dim
+    return _get_default_tile_shapes_forward(na_dim)
+def get_default_forward_tile_sizes(
+    input_tensor: Tensor,
+) -> CutlassBlackwellFmhaForwardConfigType:
+    if input_tensor.dim() != 4:
+        raise ValueError("Only 4-D tensors are supported in FMHA.")
+    q_t, kv_t = get_default_forward_tile_shapes(input_tensor)
+    assert len(q_t) == len(kv_t) == 1
+    return (q_t[0], kv_t[0])
+def get_default_backward_tile_shapes(
+    input_tensor: Tensor,
+) -> CutlassBlackwellFnaBackwardConfigType:
+    all_configs = get_all_backward_configs(input_tensor)
+    if len(all_configs) < 1:
+        device_cc = get_device_cc(input_tensor.device)
+        raise ValueError(
+            "No configs exist for this use case; Blackwell FMHA/FNA does not support it: "
+            f"{input_tensor.shape=}, {input_tensor.dtype=}, {device_cc=}."
+        )
+    return all_configs[0]
+def get_default_backward_tile_sizes(
+    input_tensor: Tensor,
+) -> CutlassBlackwellFmhaBackwardConfigType:
+    if input_tensor.dim() != 4:
+        raise ValueError("Only 4-D tensors are supported in FMHA.")
+    q_t, kv_t = get_default_backward_tile_shapes(input_tensor)
+    assert len(q_t) == len(kv_t) == 1
+    return (q_t[0], kv_t[0])
+def check_cutlass_blackwell_fna_forward_config(
+    input_tensor: Tensor,
+    q_tile_shape: Optional[DimensionType] = None,
+    kv_tile_shape: Optional[DimensionType] = None,
+) -> CutlassBlackwellFnaForwardConfigType:
+    assert input_tensor.dim() in [4, 5, 6]
+    na_dim = input_tensor.dim() - 3  # batch, heads, head_dim
+    if (q_tile_shape is None) ^ (kv_tile_shape is None):
+        raise ValueError(
+            "Please specify both q_tile_shape and kv_tile_shape, or neither one. "
+            f"Got {q_tile_shape=}, {kv_tile_shape=}."
+        )
+    if q_tile_shape is None and kv_tile_shape is None:
+        return get_default_forward_tile_shapes(input_tensor=input_tensor)
+    q_tile_shape = check_tile_shape(q_tile_shape)
+    kv_tile_shape = check_tile_shape(kv_tile_shape)
+    tile_shapes = get_all_forward_configs(input_tensor=input_tensor)
+    for q_t, kv_t in tile_shapes:
+        if q_t == q_tile_shape and kv_t == kv_tile_shape:
+            return (q_t, kv_t)  # type: ignore
+    # Fail and make suggestions
+    MAX_EXAMPLES = 3
+    examples = ""
+    for i, (q_t, kv_t) in enumerate(tile_shapes):
+        examples += f"\n  q_tile_shape={q_t}, kv_tile_shape={kv_t}"
+        if i > MAX_EXAMPLES:
+            break
+    raise ValueError(
+        f"Invalid configuration for CUTLASS Blackwell FNA-{na_dim}D. "
+        f"Q tile shape {q_tile_shape} and KV tile shape {kv_tile_shape} "
+        f"are not among the {len(tile_shapes)} configurations implementable "
+        f"with CUTLASS Blackwell FNA. "
+        "Try selecting a combination from: \n"
+        "  natten.get_configs_for_blackwell_fna(q, k, v)"
+        "\n"
+        "Here's a few examples of available combinations for your use case:\n"
+        f"{examples}"
+    )
+def check_cutlass_blackwell_fmha_forward_config(
+    input_tensor: Tensor,
+    q_tile_size: Optional[int] = None,
+    kv_tile_size: Optional[int] = None,
+) -> CutlassBlackwellFmhaForwardConfigType:
+    assert input_tensor.dim() == 4
+    if (q_tile_size is None) ^ (kv_tile_size is None):
+        raise ValueError(
+            "Please specify both q_tile_size and kv_tile_size, or neither one. "
+            f"Got {q_tile_size=}, {kv_tile_size=}."
+        )
+    if q_tile_size is None and kv_tile_size is None:
+        q_tile_size, kv_tile_size = get_default_forward_tile_sizes(
+            input_tensor=input_tensor
+        )
+        return (q_tile_size, kv_tile_size)
+    tile_sizes = get_all_fmha_forward_configs(input_tensor=input_tensor)
+    for q_t, kv_t in tile_sizes:
+        if q_t == q_tile_size and kv_t == kv_tile_size:
+            return (q_t, kv_t)
+    # Fail and make suggestions
+    MAX_EXAMPLES = 3
+    examples = ""
+    for i, (q_t, kv_t) in enumerate(tile_sizes):
+        examples += f"\n  q_tile_size={q_t}, kv_tile_size={kv_t}"
+        if i > MAX_EXAMPLES:
+            break
+    raise ValueError(
+        f"Invalid configuration for CUTLASS Blackwell FMHA. "
+        f"Q tile size {q_tile_size} and KV tile size {kv_tile_size} "
+        f"are not among the {len(tile_sizes)} configurations implementable "
+        f"with CUTLASS Blackwell FMHA. "
+        "Try selecting a combination from: \n"
+        "  natten.get_configs_for_blackwell_fmha(q, k, v)"
+        "\n"
+        "Here's a few examples of available combinations for your use case:\n"
+        f"{examples}"
+    )
+def check_cutlass_blackwell_fna_backward_config(
+    input_tensor: Tensor,
+    q_tile_shape: Optional[DimensionType] = None,
+    kv_tile_shape: Optional[DimensionType] = None,
+) -> CutlassBlackwellFnaBackwardConfigType:
+    assert input_tensor.dim() in [4, 5, 6]
+    na_dim = input_tensor.dim() - 3  # batch, heads, head_dim
+    if (q_tile_shape is None) ^ (kv_tile_shape is None):
+        raise ValueError(
+            "Please specify both q_tile_shape and kv_tile_shape, or neither one. "
+            f"Got {q_tile_shape=}, {kv_tile_shape=}."
+        )
+    if q_tile_shape is None and kv_tile_shape is None:
+        return get_default_backward_tile_shapes(input_tensor=input_tensor)
+    q_tile_shape = check_tile_shape(q_tile_shape)
+    kv_tile_shape = check_tile_shape(kv_tile_shape)
+    tile_shapes = get_all_backward_configs(input_tensor=input_tensor)
+    for q_t, kv_t in tile_shapes:
+        if q_t == q_tile_shape and kv_t == kv_tile_shape:
+            return (q_t, kv_t)  # type: ignore
+    # Fail and make suggestions
+    MAX_EXAMPLES = 3
+    examples = ""
+    for i, (q_t, kv_t) in enumerate(tile_shapes):
+        examples += f"\n  q_tile_shape={q_t}, kv_tile_shape={kv_t}"
+        if i > MAX_EXAMPLES:
+            break
+    raise ValueError(
+        f"Invalid configuration for CUTLASS Blackwell FNA-{na_dim}D Backward. "
+        f"Q tile shape {q_tile_shape} and KV tile shape {kv_tile_shape} "
+        f"are not among the {len(tile_shapes)} configurations implementable "
+        f"with CUTLASS Blackwell FNA Backward. "
+        "Try selecting a combination from: \n"
+        "  natten.get_bwd_configs_for_blackwell_fna(q, k, v)"
+        "\n"
+        "Here's a few examples of available combinations for your use case:\n"
+        f"{examples}"
+    )
+def check_cutlass_blackwell_fmha_backward_config(
+    input_tensor: Tensor,
+    q_tile_size: Optional[int] = None,
+    kv_tile_size: Optional[int] = None,
+) -> CutlassBlackwellFmhaBackwardConfigType:
+    assert input_tensor.dim() == 4
+    if (q_tile_size is None) ^ (kv_tile_size is None):
+        raise ValueError(
+            "Please specify both q_tile_size and kv_tile_size, or neither one. "
+            f"Got {q_tile_size=}, {kv_tile_size=}."
+        )
+    if q_tile_size is None and kv_tile_size is None:
+        q_tile_size, kv_tile_size = get_default_backward_tile_sizes(
+            input_tensor=input_tensor
+        )
+        return (q_tile_size, kv_tile_size)
+    tile_sizes = get_all_fmha_backward_configs(input_tensor=input_tensor)
+    for q_t, kv_t in tile_sizes:
+        if q_t == q_tile_size and kv_t == kv_tile_size:
+            return (q_t, kv_t)
+    # Fail and make suggestions
+    MAX_EXAMPLES = 3
+    examples = ""
+    for i, (q_t, kv_t) in enumerate(tile_sizes):
+        examples += f"\n  q_tile_size={q_t}, kv_tile_size={kv_t}"
+        if i > MAX_EXAMPLES:
+            break
+    raise ValueError(
+        f"Invalid configuration for CUTLASS Blackwell FMHA Backward. "
+        f"Q tile size {q_tile_size} and KV tile size {kv_tile_size} "
+        f"are not among the {len(tile_sizes)} configurations implementable "
+        f"with CUTLASS Blackwell FMHA Backward. "
+        "Try selecting a combination from: \n"
+        "  natten.get_bwd_configs_for_blackwell_fmha(q, k, v)"
+        "\n"
+        "Here's a few examples of available combinations for your use case:\n"
+        f"{examples}"
+    )

build/torch210-cu126-aarch64-linux/backends/configs/cutlass_hopper/__init__.py ADDED Viewed

	@@ -0,0 +1,522 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+from typing import List, Optional
+import torch  # noqa: F401
+from torch import Tensor
+from ...._types import (
+    CutlassHopperFmhaBackwardConfigType,
+    CutlassHopperFmhaForwardConfigType,
+    CutlassHopperFnaBackwardConfigType,
+    CutlassHopperFnaForwardConfigType,
+    DimensionType,
+    KernelSchedule,
+)
+from ....utils.checks import check_tile_shape
+from ....utils.device import get_device_cc
+DTYPE_TO_BITS = {
+    torch.float16: 16,
+    torch.bfloat16: 16,
+}
+# TODO: notes
+HOPPER_FORWARD_CONFIGS = {
+    1: {
+        16: {
+            32: [
+                (((64,), (128,)), KernelSchedule.NonPersistent),
+            ],
+            64: [
+                (((64,), (128,)), KernelSchedule.NonPersistent),
+            ],
+            128: [
+                (((128,), (128,)), KernelSchedule.WarpSpecializedCooperative),
+                (((128,), (128,)), KernelSchedule.WarpSpecializedPingpong),
+            ],
+            256: [
+                (((128,), (64,)), KernelSchedule.WarpSpecializedCooperative),
+            ],
+        },
+    },
+    2: {
+        16: {
+            32: [
+                (((8, 8), (16, 8)), KernelSchedule.NonPersistent),
+                (((8, 8), (8, 16)), KernelSchedule.NonPersistent),
+            ],
+            64: [
+                (((8, 8), (16, 8)), KernelSchedule.NonPersistent),
+                (((8, 8), (8, 16)), KernelSchedule.NonPersistent),
+            ],
+            128: [
+                (((16, 8), (16, 8)), KernelSchedule.WarpSpecializedCooperative),
+                (((16, 8), (16, 8)), KernelSchedule.WarpSpecializedPingpong),
+            ],
+            256: [
+                (((16, 8), (8, 8)), KernelSchedule.WarpSpecializedCooperative),
+                (((8, 16), (8, 8)), KernelSchedule.WarpSpecializedCooperative),
+            ],
+        },
+    },
+    3: {
+        16: {
+            32: [
+                (((4, 4, 4), (4, 4, 8)), KernelSchedule.NonPersistent),
+                (((4, 4, 4), (2, 8, 8)), KernelSchedule.NonPersistent),
+            ],
+            64: [
+                (((4, 4, 4), (4, 4, 8)), KernelSchedule.NonPersistent),
+                (((4, 4, 4), (2, 8, 8)), KernelSchedule.NonPersistent),
+            ],
+            128: [
+                (((4, 4, 8), (4, 4, 8)), KernelSchedule.WarpSpecializedCooperative),
+                (((4, 4, 8), (4, 4, 8)), KernelSchedule.WarpSpecializedPingpong),
+                (((2, 8, 8), (2, 8, 8)), KernelSchedule.WarpSpecializedCooperative),
+                (((2, 8, 8), (2, 8, 8)), KernelSchedule.WarpSpecializedPingpong),
+            ],
+            256: [
+                (((4, 4, 8), (4, 4, 4)), KernelSchedule.WarpSpecializedCooperative),
+                (((2, 8, 8), (4, 4, 4)), KernelSchedule.WarpSpecializedCooperative),
+            ],
+        },
+    },
+}
+HOPPER_BACKWARD_CONFIGS = {
+    1: {
+        16: {
+            32: [
+                ((64,), (128,)),
+                ((128,), (128,)),
+            ],
+            64: [
+                ((64,), (128,)),
+                ((128,), (128,)),
+            ],
+            128: [
+                ((64,), (128,)),
+            ],
+        },
+    },
+    2: {
+        16: {
+            32: [
+                ((8, 8), (16, 8)),
+                ((8, 8), (8, 16)),
+                ((16, 8), (16, 8)),
+                ((16, 8), (8, 16)),
+            ],
+            64: [
+                ((8, 8), (16, 8)),
+                ((8, 8), (8, 16)),
+                ((16, 8), (16, 8)),
+                ((16, 8), (8, 16)),
+            ],
+            128: [
+                ((8, 8), (16, 8)),
+                ((8, 8), (8, 16)),
+            ],
+        },
+    },
+    3: {
+        16: {
+            32: [
+                ((4, 4, 4), (4, 4, 8)),
+                ((4, 4, 4), (2, 8, 8)),
+                ((4, 4, 8), (4, 4, 8)),
+                ((4, 4, 8), (2, 8, 8)),
+            ],
+            64: [
+                ((4, 4, 4), (4, 4, 8)),
+                ((4, 4, 4), (2, 8, 8)),
+                ((4, 4, 8), (4, 4, 8)),
+                ((4, 4, 8), (2, 8, 8)),
+            ],
+            128: [
+                ((4, 4, 4), (4, 4, 8)),
+                ((4, 4, 4), (2, 8, 8)),
+                ((2, 4, 8), (2, 8, 8)),
+                ((1, 8, 8), (2, 8, 8)),
+            ],
+        },
+    },
+}
+def get_all_forward_configs(
+    input_tensor: Tensor,
+) -> List[CutlassHopperFnaForwardConfigType]:
+    assert input_tensor.dim() in [4, 5, 6]
+    na_dim = input_tensor.dim() - 3  # batch, heads, head_dim
+    dtype = input_tensor.dtype
+    dtype_bits = DTYPE_TO_BITS[dtype]
+    head_dim = input_tensor.shape[-1]
+    # if dtype not in [torch.float16, torch.bfloat16]:
+    if dtype_bits not in HOPPER_FORWARD_CONFIGS[na_dim]:  # type: ignore
+        return []
+    # if head_dim not in [32, 64, 128, 256]:
+    if head_dim not in HOPPER_FORWARD_CONFIGS[na_dim][dtype_bits]:  # type: ignore
+        return []
+    device_cc = get_device_cc(input_tensor.device)
+    if device_cc != 90:
+        return []
+    return HOPPER_FORWARD_CONFIGS[na_dim][dtype_bits][head_dim]  # type: ignore
+def get_all_backward_configs(
+    input_tensor: Tensor,
+) -> List[CutlassHopperFnaBackwardConfigType]:
+    assert input_tensor.dim() in [4, 5, 6]
+    na_dim = input_tensor.dim() - 3  # batch, heads, head_dim
+    dtype = input_tensor.dtype
+    dtype_bits = DTYPE_TO_BITS[dtype]
+    head_dim = input_tensor.shape[-1]
+    # if dtype not in [torch.float16, torch.bfloat16]:
+    if dtype_bits not in HOPPER_BACKWARD_CONFIGS[na_dim]:  # type: ignore
+        return []
+    # if head_dim not in [32, 64, 128, 256]:
+    if head_dim not in HOPPER_BACKWARD_CONFIGS[na_dim][dtype_bits]:  # type: ignore
+        return []
+    device_cc = get_device_cc(input_tensor.device)
+    if device_cc != 90:
+        return []
+    return HOPPER_BACKWARD_CONFIGS[na_dim][dtype_bits][head_dim]  # type: ignore
+# For FMHA
+def get_all_fmha_forward_configs(
+    input_tensor: Tensor,
+) -> List[CutlassHopperFmhaForwardConfigType]:
+    if input_tensor.dim() != 4:
+        raise ValueError("Only 4-D tensors are supported in FMHA.")
+    configs_multi_dim = get_all_forward_configs(input_tensor)
+    assert all(len(q_t) == len(kv_t) == 1 for (q_t, kv_t), _ in configs_multi_dim)
+    configs_fmha = [
+        ((q_t[0], kv_t[0]), sched) for (q_t, kv_t), sched in configs_multi_dim
+    ]
+    return configs_fmha
+def get_all_fmha_backward_configs(
+    input_tensor: Tensor,
+) -> List[CutlassHopperFmhaBackwardConfigType]:
+    if input_tensor.dim() != 4:
+        raise ValueError("Only 4-D tensors are supported in FMHA.")
+    configs_multi_dim = get_all_backward_configs(input_tensor)
+    assert all(len(q_t) == len(kv_t) == 1 for q_t, kv_t in configs_multi_dim)
+    configs_fmha = [(q_t[0], kv_t[0]) for q_t, kv_t in configs_multi_dim]
+    return configs_fmha
+def get_default_forward_config(
+    input_tensor: Tensor,
+) -> CutlassHopperFnaForwardConfigType:
+    all_configs = get_all_forward_configs(input_tensor)
+    if len(all_configs) < 1:
+        device_cc = get_device_cc(input_tensor.device)
+        raise ValueError(
+            "No configs exist for this use case; Hopper FMHA/FNA does not support it: "
+            f"{input_tensor.shape=}, {input_tensor.dtype=}, {device_cc=}."
+        )
+    return all_configs[0]
+def get_default_backward_config(
+    input_tensor: Tensor,
+) -> CutlassHopperFnaBackwardConfigType:
+    all_configs = get_all_backward_configs(input_tensor)
+    if len(all_configs) < 1:
+        device_cc = get_device_cc(input_tensor.device)
+        raise ValueError(
+            "No configs exist for this use case; Hopper FMHA/FNA does not support it: "
+            f"{input_tensor.shape=}, {input_tensor.dtype=}, {device_cc=}."
+        )
+    return all_configs[0]
+def get_default_fmha_forward_config(
+    input_tensor: Tensor,
+) -> CutlassHopperFmhaForwardConfigType:
+    if input_tensor.dim() != 4:
+        raise ValueError("Only 4-D tensors are supported in FMHA.")
+    (q_t, kv_t), sched = get_default_forward_config(input_tensor)
+    assert len(q_t) == len(kv_t) == 1
+    return (q_t[0], kv_t[0]), sched
+def get_default_fmha_backward_config(
+    input_tensor: Tensor,
+) -> CutlassHopperFmhaBackwardConfigType:
+    if input_tensor.dim() != 4:
+        raise ValueError("Only 4-D tensors are supported in FMHA.")
+    q_t, kv_t = get_default_backward_config(input_tensor)
+    assert len(q_t) == len(kv_t) == 1
+    return q_t[0], kv_t[0]
+def check_cutlass_hopper_fna_forward_config(
+    input_tensor: Tensor,
+    q_tile_shape: Optional[DimensionType] = None,
+    kv_tile_shape: Optional[DimensionType] = None,
+    kernel_schedule: Optional[KernelSchedule] = None,
+) -> CutlassHopperFnaForwardConfigType:
+    assert input_tensor.dim() in [4, 5, 6]
+    na_dim = input_tensor.dim() - 3  # batch, heads, head_dim
+    if (q_tile_shape is None) ^ (kv_tile_shape is None):
+        raise ValueError(
+            "Please specify both q_tile_shape and kv_tile_shape, or neither one. "
+            f"Got {q_tile_shape=}, {kv_tile_shape=}."
+        )
+    (default_q_tile_shape, default_kv_tile_shape), default_sched = (
+        get_default_forward_config(input_tensor=input_tensor)
+    )
+    if q_tile_shape is None and kv_tile_shape is None and kernel_schedule is None:
+        return (default_q_tile_shape, default_kv_tile_shape), default_sched  # type: ignore[return-value]
+    elif q_tile_shape is None and kv_tile_shape is None:
+        q_tile_shape = default_q_tile_shape
+        kv_tile_shape = default_kv_tile_shape
+    q_tile_shape = check_tile_shape(q_tile_shape)
+    kv_tile_shape = check_tile_shape(kv_tile_shape)
+    configs = get_all_forward_configs(input_tensor=input_tensor)
+    for (q_t, kv_t), sched in configs:
+        if (
+            q_t == q_tile_shape
+            and kv_t == kv_tile_shape
+            and (kernel_schedule is None or sched == kernel_schedule)
+        ):
+            return (q_t, kv_t), sched  # type: ignore
+    # Fail and make suggestions
+    MAX_EXAMPLES = 3
+    examples = ""
+    for i, ((q_t, kv_t), sched) in enumerate(configs):
+        examples += f"\n  q_tile_shape={q_t}, kv_tile_shape={kv_t}, schedule={sched}"
+        if i > MAX_EXAMPLES:
+            break
+    raise ValueError(
+        f"Invalid configuration for CUTLASS Hopper FNA-{na_dim}D. "
+        f"Q tile shape {q_tile_shape}, KV tile shape {kv_tile_shape}, and schedule {kernel_schedule} "
+        f"are not among the {len(configs)} configurations implementable "
+        f"with CUTLASS Hopper FNA. "
+        "Try selecting a combination from: \n"
+        "  natten.get_configs_for_hopper_fna(q, k, v)"
+        "\n"
+        "Here's a few examples of available combinations for your use case:\n"
+        f"{examples}"
+    )
+def check_cutlass_hopper_fna_backward_config(
+    input_tensor: Tensor,
+    q_tile_shape: Optional[DimensionType] = None,
+    kv_tile_shape: Optional[DimensionType] = None,
+) -> CutlassHopperFnaBackwardConfigType:
+    assert input_tensor.dim() in [4, 5, 6]
+    na_dim = input_tensor.dim() - 3  # batch, heads, head_dim
+    if (q_tile_shape is None) ^ (kv_tile_shape is None):
+        raise ValueError(
+            "Please specify both q_tile_shape and kv_tile_shape, or neither one. "
+            f"Got {q_tile_shape=}, {kv_tile_shape=}."
+        )
+    default_q_tile_shape, default_kv_tile_shape = get_default_backward_config(
+        input_tensor=input_tensor
+    )
+    if q_tile_shape is None and kv_tile_shape is None:
+        return default_q_tile_shape, default_kv_tile_shape  # type: ignore[return-value]
+    elif q_tile_shape is None and kv_tile_shape is None:
+        q_tile_shape = default_q_tile_shape
+        kv_tile_shape = default_kv_tile_shape
+    q_tile_shape = check_tile_shape(q_tile_shape)
+    kv_tile_shape = check_tile_shape(kv_tile_shape)
+    configs = get_all_backward_configs(input_tensor=input_tensor)
+    for q_t, kv_t in configs:
+        if q_t == q_tile_shape and kv_t == kv_tile_shape:
+            return q_t, kv_t  # type: ignore
+    # Fail and make suggestions
+    MAX_EXAMPLES = 3
+    examples = ""
+    for i, (q_t, kv_t) in enumerate(configs):
+        examples += f"\n  q_tile_shape={q_t}, kv_tile_shape={kv_t}"
+        if i > MAX_EXAMPLES:
+            break
+    raise ValueError(
+        f"Invalid configuration for CUTLASS Hopper FNA-{na_dim}D Backward. "
+        f"Q tile shape {q_tile_shape} and KV tile shape {kv_tile_shape} "
+        f"are not among the {len(configs)} configurations implementable "
+        f"with CUTLASS Hopper FNA Backward. "
+        "Try selecting a combination from: \n"
+        "  natten.get_bwd_configs_for_hopper_fna(q, k, v)"
+        "\n"
+        "Here's a few examples of available combinations for your use case:\n"
+        f"{examples}"
+    )
+def check_cutlass_hopper_fmha_forward_config(
+    input_tensor: Tensor,
+    q_tile_size: Optional[int] = None,
+    kv_tile_size: Optional[int] = None,
+    kernel_schedule: Optional[KernelSchedule] = None,
+) -> CutlassHopperFmhaForwardConfigType:
+    assert input_tensor.dim() == 4
+    if (q_tile_size is None) ^ (kv_tile_size is None):
+        raise ValueError(
+            "Please specify both q_tile_size and kv_tile_size, or neither one. "
+            f"Got {q_tile_size=}, {kv_tile_size=}."
+        )
+    (default_q_tile_size, default_kv_tile_size), default_sched = (
+        get_default_fmha_forward_config(input_tensor=input_tensor)
+    )
+    if q_tile_size is None and kv_tile_size is None and kernel_schedule is None:
+        return (default_q_tile_size, default_kv_tile_size), default_sched
+    elif q_tile_size is None and kv_tile_size is None:
+        q_tile_size = default_q_tile_size
+        kv_tile_size = default_kv_tile_size
+    configs = get_all_fmha_forward_configs(input_tensor=input_tensor)
+    for (q_t, kv_t), sched in configs:
+        if (
+            q_t == q_tile_size
+            and kv_t == kv_tile_size
+            and (kernel_schedule is None or sched == kernel_schedule)
+        ):
+            return (q_t, kv_t), sched  # type: ignore
+    # Fail and make suggestions
+    MAX_EXAMPLES = 3
+    examples = ""
+    for i, ((q_t, kv_t), sched) in enumerate(configs):
+        examples += f"\n  q_tile_size={q_t}, kv_tile_size={kv_t}, schedule={sched}"
+        if i > MAX_EXAMPLES:
+            break
+    raise ValueError(
+        f"Invalid configuration for CUTLASS Hopper FMHA. "
+        f"Q tile size {q_tile_size}, KV tile size {kv_tile_size}, and schedule {kernel_schedule} "
+        f"are not among the {len(configs)} configurations implementable "
+        f"with CUTLASS Hopper FMHA. "
+        "Try selecting a combination from: \n"
+        "  natten.get_configs_for_hopper_fmha(q, k, v)"
+        "\n"
+        "Here's a few examples of available combinations for your use case:\n"
+        f"{examples}"
+    )
+def check_cutlass_hopper_fmha_backward_config(
+    input_tensor: Tensor,
+    q_tile_size: Optional[int] = None,
+    kv_tile_size: Optional[int] = None,
+) -> CutlassHopperFmhaBackwardConfigType:
+    assert input_tensor.dim() == 4
+    if (q_tile_size is None) ^ (kv_tile_size is None):
+        raise ValueError(
+            "Please specify both q_tile_size and kv_tile_size, or neither one. "
+            f"Got {q_tile_size=}, {kv_tile_size=}."
+        )
+    default_q_tile_size, default_kv_tile_size = get_default_fmha_backward_config(
+        input_tensor=input_tensor
+    )
+    if q_tile_size is None and kv_tile_size is None:
+        return default_q_tile_size, default_kv_tile_size
+    elif q_tile_size is None and kv_tile_size is None:
+        q_tile_size = default_q_tile_size
+        kv_tile_size = default_kv_tile_size
+    configs = get_all_fmha_backward_configs(input_tensor=input_tensor)
+    for q_t, kv_t in configs:
+        if q_t == q_tile_size and kv_t == kv_tile_size:
+            return q_t, kv_t  # type: ignore
+    # Fail and make suggestions
+    MAX_EXAMPLES = 3
+    examples = ""
+    for i, (q_t, kv_t) in enumerate(configs):
+        examples += f"\n  q_tile_size={q_t}, kv_tile_size={kv_t}"
+        if i > MAX_EXAMPLES:
+            break
+    raise ValueError(
+        f"Invalid configuration for CUTLASS Hopper FMHA Backward. "
+        f"Q tile size {q_tile_size} and KV tile size {kv_tile_size} "
+        f"are not among the {len(configs)} configurations implementable "
+        f"with CUTLASS Hopper FMHA Backward. "
+        "Try selecting a combination from: \n"
+        "  natten.get_bwd_configs_for_hopper_fmha(q, k, v)"
+        "\n"
+        "Here's a few examples of available combinations for your use case:\n"
+        f"{examples}"
+    )

build/torch210-cu126-aarch64-linux/backends/configs/flex/__init__.py ADDED Viewed

	@@ -0,0 +1,210 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+from typing import List, Optional
+import torch  # noqa: F401
+from torch import Tensor
+from ...._types import (
+    DimensionType,
+    FlexFmhaForwardConfigType,
+    FlexFnaForwardConfigType,
+)
+from ....utils.checks import check_tile_shape
+from ....utils.device import get_device_cc
+# TODO: add more tile sizes/shapes
+# TODO: add backprop tile sizes/shapes
+# Only doing 64 x 64 for now, since it's the one that successfully compiles across devices and
+# use cases without running into compile errors (i.e. shmem over-subscription)
+# Once Flex with compilation actually starts working as expected and is out of prototype, we can
+# add in more tile sizes/shapes and condition them on arch / use case, like we do for CUTLASS FNA.
+FLEX_FORWARD_TILE_SHAPES = {
+    1: [
+        # ((128, ), (128, )),
+        ((64,), (64,)),
+    ],
+    2: [
+        # ((8, 16), (8, 16)),
+        ((8, 8), (8, 8)),
+        ((4, 16), (4, 16)),
+        ((4, 16), (8, 8)),
+    ],
+    3: [
+        # ((4, 4, 8), (4, 4, 8)),
+        ((4, 4, 4), (4, 4, 4)),
+        ((2, 4, 8), (2, 4, 8)),
+        ((2, 4, 8), (4, 4, 4)),
+    ],
+}
+def _get_default_tile_shapes_forward(
+    na_dim: int,
+) -> FlexFnaForwardConfigType:
+    assert na_dim in [1, 2, 3]
+    if na_dim == 1:
+        return ((64,), (64,))
+    if na_dim == 2:
+        return ((8, 8), (8, 8))
+    if na_dim == 3:
+        return ((4, 4, 4), (4, 4, 4))
+    raise NotImplementedError()
+def get_all_tile_shapes_forward(
+    input_tensor: Tensor,
+) -> List[FlexFnaForwardConfigType]:
+    assert input_tensor.dim() in [4, 5, 6]
+    na_dim = input_tensor.dim() - 3  # batch, heads, head_dim
+    return FLEX_FORWARD_TILE_SHAPES[na_dim]  # type: ignore
+# For FMHA
+def get_all_tile_sizes_forward(input_tensor: Tensor) -> List[FlexFmhaForwardConfigType]:
+    if input_tensor.dim() != 4:
+        raise ValueError("Only 4-D tensors are supported in FMHA.")
+    tile_shapes = get_all_tile_shapes_forward(input_tensor)
+    assert all(len(q_t) == len(kv_t) == 1 for q_t, kv_t in tile_shapes)
+    tile_sizes = [(q_t[0], kv_t[0]) for q_t, kv_t in tile_shapes]
+    return tile_sizes
+def get_default_forward_tile_shapes(input_tensor: Tensor) -> FlexFnaForwardConfigType:
+    assert input_tensor.dim() in [4, 5, 6]
+    na_dim = input_tensor.dim() - 3  # batch, heads, head_dim
+    return _get_default_tile_shapes_forward(na_dim)
+def get_default_forward_tile_sizes(input_tensor: Tensor) -> FlexFmhaForwardConfigType:
+    if input_tensor.dim() != 4:
+        raise ValueError("Only 4-D tensors are supported in FMHA.")
+    q_t, kv_t = get_default_forward_tile_shapes(input_tensor)
+    assert len(q_t) == len(kv_t) == 1
+    return (q_t[0], kv_t[0])
+def check_flex_fna_forward_config(
+    input_tensor: Tensor,
+    q_tile_shape: Optional[DimensionType] = None,
+    kv_tile_shape: Optional[DimensionType] = None,
+) -> FlexFnaForwardConfigType:
+    assert input_tensor.dim() in [4, 5, 6]
+    na_dim = input_tensor.dim() - 3  # batch, heads, head_dim
+    if (q_tile_shape is None) ^ (kv_tile_shape is None):
+        raise ValueError(
+            "Please specify both q_tile_shape and kv_tile_shape, or neither one. "
+            f"Got {q_tile_shape=}, {kv_tile_shape=}."
+        )
+    if q_tile_shape is None and kv_tile_shape is None:
+        return get_default_forward_tile_shapes(input_tensor=input_tensor)
+    q_tile_shape = check_tile_shape(q_tile_shape)
+    kv_tile_shape = check_tile_shape(kv_tile_shape)
+    tile_shapes = get_all_tile_shapes_forward(input_tensor=input_tensor)
+    for q_t, kv_t in tile_shapes:
+        if q_t == q_tile_shape and kv_t == kv_tile_shape:
+            return (q_t, kv_t)  # type: ignore
+    # Fail and make suggestions
+    device_cc = get_device_cc(input_tensor.device)
+    MAX_EXAMPLES = 3
+    examples = ""
+    for i, (q_t, kv_t) in enumerate(tile_shapes):
+        examples += f"\n  q_tile_shape={q_t}, kv_tile_shape={kv_t}"
+        if i > MAX_EXAMPLES:
+            break
+    raise ValueError(
+        f"Invalid configuration for Flex FNA-{na_dim}D. "
+        f"Q tile shape {q_tile_shape} and KV tile shape {kv_tile_shape} "
+        f"are not among the {len(tile_shapes)} configurations implementable "
+        f"with Flex FNA for SM{device_cc}. "
+        "Try selecting a combination from: \n"
+        "  natten.get_configs_for_flex_fna(q, k, v)"
+        "\n"
+        "Here's a few examples of available combinations for your use case:\n"
+        f"{examples}"
+    )
+def check_flex_fmha_forward_config(
+    input_tensor: Tensor,
+    q_tile_size: Optional[int] = None,
+    kv_tile_size: Optional[int] = None,
+) -> FlexFmhaForwardConfigType:
+    assert input_tensor.dim() == 4
+    if (q_tile_size is None) ^ (kv_tile_size is None):
+        raise ValueError(
+            "Please specify both q_tile_size and kv_tile_size, or neither one. "
+            f"Got {q_tile_size=}, {kv_tile_size=}."
+        )
+    if q_tile_size is None and kv_tile_size is None:
+        q_tile_size, kv_tile_size = get_default_forward_tile_sizes(
+            input_tensor=input_tensor
+        )
+        return (q_tile_size, kv_tile_size)
+    tile_sizes = get_all_tile_sizes_forward(input_tensor=input_tensor)
+    for q_t, kv_t in tile_sizes:
+        if q_t == q_tile_size and kv_t == kv_tile_size:
+            return (q_t, kv_t)
+    # Fail and make suggestions
+    device_cc = get_device_cc(input_tensor.device)
+    MAX_EXAMPLES = 3
+    examples = ""
+    for i, (q_t, kv_t) in enumerate(tile_sizes):
+        examples += f"\n  q_tile_size={q_t}, kv_tile_size={kv_t}"
+        if i > MAX_EXAMPLES:
+            break
+    raise ValueError(
+        f"Invalid configuration for Flex FMHA. "
+        f"Q tile size {q_tile_size} and KV tile size {kv_tile_size} "
+        f"are not among the {len(tile_sizes)} configurations implementable "
+        f"with Flex FMHA for SM{device_cc}. "
+        "Try selecting a combination from: \n"
+        "  natten.get_configs_for_flex_fmha(q, k, v)\n"
+        "\n"
+        "Here's a few examples of available combinations for your use case:\n"
+        f"{examples}"
+    )

build/torch210-cu126-aarch64-linux/backends/flex.py ADDED Viewed

	@@ -0,0 +1,799 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+import functools
+import math
+import time
+from typing import Callable, Optional, Tuple, Union
+import torch
+from torch import BoolTensor, IntTensor, Tensor
+from torch.nn.attention.flex_attention import (
+    BlockMask,
+    create_block_mask,
+    flex_attention,
+)
+from ..backends.configs.checks import (  # noqa: F401
+    _FLEX_COMPILE_SUPPORTED,
+    _FLEX_SUPPORTED,
+    can_run_flex_attention,
+)
+from ..backends.configs.flex import (
+    check_flex_fmha_forward_config,
+    check_flex_fna_forward_config,
+)
+from ..token_permute import token_permute_operation, token_unpermute_operation
+from .._types import (
+    CausalArg1DTypeOrDed,
+    CausalArg2DTypeOrDed,
+    CausalArg3DTypeOrDed,
+    CausalArgType,
+    CausalArgTypeOrDed,
+    Dimension1DType,
+    Dimension1DTypeOrDed,
+    Dimension2DType,
+    Dimension2DTypeOrDed,
+    Dimension3DType,
+    Dimension3DTypeOrDed,
+    DimensionType,
+    DimensionTypeOrDed,
+)
+from ..utils import log
+from ..utils.checks import (
+    check_all_args,
+    check_args_against_input,
+    check_input_size_arg,
+    fmha_tensor_checks,
+    na_tensor_checks,
+    varlen_tensor_checks,
+)
+from ..utils.environment import is_torch_compiling
+logger = log.get_logger(__name__)
+def get_flex_attention_fn(
+    torch_compile: bool, torch_compile_args: Optional[dict] = None
+) -> Callable:
+    if not torch_compile:
+        return flex_attention
+    additional_args = torch_compile_args or {}
+    additional_args["dynamic"] = False
+    return torch.compile(flex_attention, **additional_args)
+def _run_flex_attn(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    block_mask: BlockMask,
+    scale: float,
+    torch_compile: bool,
+    q_tile_size: Optional[int] = None,
+    kv_tile_size: Optional[int] = None,
+    torch_compile_args: Optional[dict] = None,
+) -> Tuple[Tensor, Tensor]:
+    # We may need to override the default flex config.
+    # Default ones are not guaranteed to work out of the box across architectures.
+    # Some oversubscribe shmem even on the B200!
+    torch_compile_args = {}
+    # Disable flex decoding path
+    kernel_options = {
+        "FORCE_USE_FLEX_ATTENTION": True,
+    }
+    if q_tile_size is not None and torch_compile:
+        kv_tile_size = kv_tile_size or q_tile_size
+        # Have to auto-tune, otherwise torch will only allow the default config.
+        torch_compile_args["mode"] = "max-autotune-no-cudagraphs"
+        kernel_options["SPARSE_Q_BLOCK_SIZE"] = q_tile_size  # type: ignore[assignment]
+        kernel_options["SPARSE_KV_BLOCK_SIZE"] = kv_tile_size  # type: ignore[assignment]
+        kernel_options["BLOCK_M"] = q_tile_size  # type: ignore[assignment]
+        kernel_options["BLOCK_N"] = kv_tile_size  # type: ignore[assignment]
+    flex_fn = get_flex_attention_fn(
+        torch_compile=torch_compile, torch_compile_args=torch_compile_args
+    )
+    # tensors are BHSD here
+    is_gqa = q.shape[1] != k.shape[1]
+    return flex_fn(
+        q,
+        k,
+        v,
+        block_mask=block_mask,
+        return_lse=True,
+        scale=scale,
+        kernel_options=kernel_options,
+        enable_gqa=is_gqa,
+    )
+def run_flex_attn(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    block_mask: BlockMask,
+    scale: float,
+    torch_compile: bool,
+    torch_compile_args: Optional[dict] = None,
+    q_tile_size: Optional[int] = None,
+    kv_tile_size: Optional[int] = None,
+) -> Tuple[Tensor, Tensor]:
+    if q_tile_size is not None and kv_tile_size is not None:
+        return _run_flex_attn(
+            q,
+            k,
+            v,
+            block_mask=block_mask,
+            scale=scale,
+            torch_compile=torch_compile,
+            q_tile_size=q_tile_size,
+            kv_tile_size=kv_tile_size,
+            torch_compile_args=torch_compile_args,
+        )
+    # Use smallest tile size combo to try and evade shmem oversubscription
+    # The defaults just fail very frequently.
+    return _run_flex_attn(
+        q,
+        k,
+        v,
+        block_mask=block_mask,
+        scale=scale,
+        torch_compile=torch_compile,
+        q_tile_size=64,
+        kv_tile_size=64,
+        torch_compile_args=torch_compile_args,
+    )
+def flex_fmha(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    is_causal: bool = False,
+    scale: Optional[float] = None,
+    q_tile_size: Optional[int] = None,
+    kv_tile_size: Optional[int] = None,
+    torch_compile: bool = False,
+    return_lse: bool = False,
+    # varlen parameters
+    cumulative_seqlen_Q: Optional[Tensor] = None,
+    cumulative_seqlen_KV: Optional[Tensor] = None,
+    max_seqlen_Q: int = 0,
+    max_seqlen_KV: int = 0,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    fmha_tensor_checks(
+        query,
+        key,
+        value,
+        must_match_head_dims=True,
+        supports_gqa_mqa=True,
+        backend_name="Flex FMHA",
+    )
+    (
+        cumulative_seqlen_Q,
+        cumulative_seqlen_KV,
+        max_seqlen_Q,
+        max_seqlen_KV,
+    ) = varlen_tensor_checks(
+        query=query,
+        key=key,
+        value=value,
+        cumulative_seqlen_Q=cumulative_seqlen_Q,
+        cumulative_seqlen_KV=cumulative_seqlen_KV,
+        max_seqlen_Q=max_seqlen_Q,
+        max_seqlen_KV=max_seqlen_KV,
+    )
+    is_varlen = cumulative_seqlen_Q is not None
+    assert can_run_flex_attention(
+        query,
+        key,
+        value,
+        is_causal=is_causal,
+        is_varlen=is_varlen,
+        torch_compile=torch_compile,
+        raise_error=True,
+    )
+    q_tile_size, kv_tile_size = check_flex_fmha_forward_config(
+        input_tensor=query,
+        q_tile_size=q_tile_size,
+        kv_tile_size=kv_tile_size,
+    )
+    scale = scale or query.shape[-1] ** -0.5
+    batch_size, seqlen_q, num_heads, head_dim = query.shape
+    _, seqlen_kv, num_heads_kv, head_dim_v = value.shape
+    # Flex and torch attention use heads first layout
+    query_ = query.reshape(batch_size, seqlen_q, num_heads, head_dim).transpose(1, 2)
+    key_ = key.reshape(batch_size, seqlen_kv, num_heads_kv, head_dim_v).transpose(1, 2)
+    value_ = value.reshape(batch_size, seqlen_kv, num_heads_kv, head_dim_v).transpose(
+        1, 2
+    )
+    out_, lse_ = run_flex_attn(
+        query_,
+        key_,
+        value_,
+        block_mask=None,  # type: ignore[arg-type]
+        scale=scale,
+        torch_compile=torch_compile,
+        q_tile_size=q_tile_size,
+        kv_tile_size=kv_tile_size,
+    )
+    out = out_.transpose(1, 2).reshape(batch_size, seqlen_q, num_heads, head_dim_v)
+    lse = lse_.transpose(1, 2).reshape(batch_size, seqlen_q, num_heads)
+    if return_lse:
+        return out, lse
+    return out
+# TODO: move me elsewhere?
+def idx2crd(index, shape) -> tuple:
+    rank = len(shape)
+    coord = []
+    residual = index
+    for i in range(rank - 1, -1, -1):
+        coord.append(residual % shape[i])
+        residual = residual // shape[i]
+    # assert residual == 0
+    return tuple(coord[::-1])
+def get_na_flex_mask(
+    device: str,
+    na_dim: int,
+    qkv_shape: DimensionType,
+    kernel_size: DimensionType,
+    stride: DimensionType,
+    dilation: DimensionType,
+    is_causal: CausalArgType,
+    q_tile_shape: Optional[DimensionType] = None,
+    kv_tile_shape: Optional[DimensionType] = None,
+    q_shape: Optional[DimensionType] = None,
+    kv_shape: Optional[DimensionType] = None,
+    torch_compile: bool = False,
+):
+    num_dilation_groups = math.prod(dilation)
+    if not is_torch_compiling():
+        flex_mask_start_time = time.perf_counter()
+    do_token_permute = q_tile_shape is not None and kv_tile_shape is not None
+    if do_token_permute:
+        if q_tile_shape is None or kv_tile_shape is None:
+            raise ValueError(
+                "Please specify Q and KV tile shapes for multi dimensional tiling. "
+                f"Got {q_tile_shape=}, {kv_tile_shape=}."
+            )
+        if q_shape is None or kv_shape is None:
+            raise ValueError(
+                "Please specify q_shape and kv_shape for multi dimensional tiling."
+            )
+        if len(q_tile_shape) != na_dim or len(kv_tile_shape) != na_dim:
+            raise ValueError(
+                "Q and KV tile shapes must match the number of dimensions in the "
+                f"token layout ({na_dim}, got {q_tile_shape=}, {kv_tile_shape=}."
+            )
+        if any(x % t != 0 for x, t in zip(q_shape, q_tile_shape)):
+            raise ValueError(
+                "Input must be divisible by Q tile shape, but got "
+                f"{q_shape=}, {q_tile_shape=}."
+            )
+        if any(x % t != 0 for x, t in zip(kv_shape, kv_tile_shape)):
+            raise ValueError(
+                "Input must be divisible by KV tile shape, but got "
+                f"{kv_shape=}, {kv_tile_shape=}."
+            )
+        q_rest_shape = tuple(x // t for x, t in zip(q_shape, q_tile_shape))
+        kv_rest_shape = tuple(x // t for x, t in zip(kv_shape, kv_tile_shape))
+    def single_dim_tiling_mask(
+        b: IntTensor,
+        h: IntTensor,
+        q_idx: IntTensor,
+        kv_idx: IntTensor,
+        qkv_shape,
+        kernel_size,
+        stride,
+        dilation,
+        is_causal,
+    ) -> BoolTensor:
+        # Reconstruct global Q and KV coordinates
+        q_crd = idx2crd(q_idx, qkv_shape)
+        kv_crd = idx2crd(kv_idx, qkv_shape)
+        # Coordinates within dilation group
+        q_crd_di = tuple(x // d for x, d in zip(q_crd, dilation))
+        kv_crd_di = tuple(x // d for x, d in zip(kv_crd, dilation))
+        # Dilation group coordinates
+        q_dilation_group_crd = tuple(x % d for x, d in zip(q_crd, dilation))
+        kv_dilation_group_crd = tuple(x % d for x, d in zip(kv_crd, dilation))
+        # Fixup input shape according to dilation group
+        dilation_group_padding = tuple(
+            1 - ((dg + (d - (x % d))) // d)
+            for dg, d, x in zip(q_dilation_group_crd, dilation, qkv_shape)
+        )
+        qkv_shape_corrected = tuple(
+            (x // d) + p for p, d, x in zip(dilation_group_padding, dilation, qkv_shape)
+        )
+        # Window size left and right (non-causal only)
+        window_size_left = tuple(w // 2 for w in kernel_size)
+        window_size_right = tuple(w // 2 + (w % 2 - 1) for w in kernel_size)
+        masks = []
+        for i in range(na_dim):
+            if is_causal[i]:
+                # Leader is the last (right-most) query in the stride group.
+                stride_group_leader = torch.min(
+                    (q_crd_di[i] // stride[i]) * stride[i] + stride[i] - 1,
+                    qkv_shape_corrected[i] - 1,
+                )
+                mask = (
+                    (
+                        q_crd_di[i] - kv_crd_di[i] >= 0
+                    )  # window still ends at query index
+                    & (stride_group_leader - kv_crd_di[i] < kernel_size[i])
+                    & (q_dilation_group_crd[i] == kv_dilation_group_crd[i])
+                )
+            else:
+                # Leader is the center-most query in the stride group.
+                # If stride is even, choose the right hand side center query.
+                stride_group_leader = torch.min(
+                    (q_crd_di[i] // stride[i]) * stride[i] + (stride[i] // 2),
+                    qkv_shape_corrected[i] - 1,
+                )
+                window_center = stride_group_leader.clamp(
+                    window_size_left[i] * torch.ones_like(qkv_shape_corrected[i]),
+                    qkv_shape_corrected[i] - 1 - window_size_right[i],
+                )
+                w0 = window_center - kv_crd_di[i]
+                w1 = kv_crd_di[i] - window_center
+                mask = (
+                    ((0 <= w0) & (w0 <= window_size_left[i]))
+                    | ((0 <= w1) & (w1 <= window_size_right[i]))
+                ) & (q_dilation_group_crd[i] == kv_dilation_group_crd[i])
+            masks.append(mask)
+        return functools.reduce(lambda x, y: x & y, masks)  # type: ignore
+    def multi_dim_tiling_mask(
+        b: IntTensor,
+        h: IntTensor,
+        q_idx: IntTensor,
+        kv_idx: IntTensor,
+        q_tile_size: int,
+        kv_tile_size: int,
+        q_tile_shape,
+        kv_tile_shape,
+        qkv_shape,
+        kernel_size,
+        stride,
+        dilation,
+        is_causal,
+    ) -> BoolTensor:
+        # Reconstruct global Q and KV coordinates
+        q_tile_idx = q_idx // q_tile_size
+        kv_tile_idx = kv_idx // kv_tile_size
+        q_tile_offset = q_idx % q_tile_size
+        kv_tile_offset = kv_idx % q_tile_size
+        q_tile_coord = idx2crd(q_tile_idx, q_rest_shape)
+        kv_tile_coord = idx2crd(kv_tile_idx, kv_rest_shape)
+        q_tile_offset_coord = idx2crd(q_tile_offset, q_tile_shape)
+        kv_tile_offset_coord = idx2crd(kv_tile_offset, kv_tile_shape)
+        q_crd = tuple(
+            tile_crd * tile_sz + tile_off
+            for tile_crd, tile_sz, tile_off in zip(
+                q_tile_coord, q_tile_shape, q_tile_offset_coord
+            )
+        )
+        kv_crd = tuple(
+            tile_crd * tile_sz + tile_off
+            for tile_crd, tile_sz, tile_off in zip(
+                kv_tile_coord, kv_tile_shape, kv_tile_offset_coord
+            )
+        )
+        # Dilation group coordinates
+        # b_actual = b // num_dilation_groups
+        dilation_group_idx = b % num_dilation_groups
+        dilation_group_crd = idx2crd(dilation_group_idx, dilation)
+        # Fixup input shape according to dilation group
+        dilation_group_padding = tuple(
+            1 - ((dg + (d - (x % d))) // d)
+            for dg, d, x in zip(dilation_group_crd, dilation, qkv_shape)
+        )
+        qkv_shape_corrected = tuple(
+            (x // d) + p for p, d, x in zip(dilation_group_padding, dilation, qkv_shape)
+        )
+        # Window size left and right (non-causal only)
+        window_size_left = tuple(w // 2 for w in kernel_size)
+        window_size_right = tuple(w // 2 + (w % 2 - 1) for w in kernel_size)
+        masks = []
+        for i in range(na_dim):
+            if is_causal[i]:
+                # Leader is the last (right-most) query in the stride group.
+                stride_group_leader = torch.min(
+                    (q_crd[i] // stride[i]) * stride[i] + stride[i] - 1,
+                    qkv_shape_corrected[i] - 1,
+                )
+                mask = (
+                    q_crd[i] - kv_crd[i] >= 0
+                ) & (  # window still ends at query index
+                    stride_group_leader - kv_crd[i] < kernel_size[i]
+                )
+            else:
+                # Leader is the center-most query in the stride group.
+                # If stride is even, choose the right hand side center query.
+                stride_group_leader = torch.min(
+                    (q_crd[i] // stride[i]) * stride[i] + (stride[i] // 2),
+                    qkv_shape_corrected[i] - 1,
+                )
+                window_center = stride_group_leader.clamp(
+                    window_size_left[i] * torch.ones_like(qkv_shape_corrected[i]),
+                    qkv_shape_corrected[i] - 1 - window_size_right[i],
+                )
+                w0 = window_center - kv_crd[i]
+                w1 = kv_crd[i] - window_center
+                mask = ((0 <= w0) & (w0 <= window_size_left[i])) | (
+                    (0 <= w1) & (w1 <= window_size_right[i])
+                )
+            masks.append(mask)
+        return functools.reduce(lambda x, y: x & y, masks)  # type: ignore
+    mask_mod = None
+    seq_length_q = seq_length_kv = math.prod(qkv_shape)
+    q_tile_size, kv_tile_size = 64, 64
+    if do_token_permute:
+        assert q_shape is not None
+        assert kv_shape is not None
+        assert q_tile_shape is not None
+        assert kv_tile_shape is not None
+        seq_length_q = math.prod(q_shape)
+        seq_length_kv = math.prod(kv_shape)
+        q_tile_size, kv_tile_size = math.prod(q_tile_shape), math.prod(kv_tile_shape)
+        mask_mod = functools.partial(
+            multi_dim_tiling_mask,
+            q_tile_size=q_tile_size,
+            kv_tile_size=kv_tile_size,
+            q_tile_shape=q_tile_shape,
+            kv_tile_shape=kv_tile_shape,
+            qkv_shape=qkv_shape,
+            kernel_size=kernel_size,
+            stride=stride,
+            dilation=dilation,
+            is_causal=is_causal,
+        )
+    else:
+        mask_mod = functools.partial(
+            single_dim_tiling_mask,
+            qkv_shape=qkv_shape,
+            kernel_size=kernel_size,
+            stride=stride,
+            dilation=dilation,
+            is_causal=is_causal,
+        )
+    block_mask = create_block_mask(
+        mask_mod,
+        B=None,
+        H=None,
+        Q_LEN=seq_length_q,
+        KV_LEN=seq_length_kv,
+        _compile=torch_compile,
+        BLOCK_SIZE=(q_tile_size, kv_tile_size),
+        device=device,
+    )
+    if not is_torch_compiling():
+        flex_mask_end_time = time.perf_counter()
+        flex_mask_time = flex_mask_end_time - flex_mask_start_time
+        logger.debug(
+            f"Flex Attention block mask ({torch_compile=}) created in {flex_mask_time:.2f} seconds."
+        )
+    return block_mask
+def flex_fna_generic(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    kernel_size: DimensionTypeOrDed,
+    stride: DimensionTypeOrDed = 1,
+    dilation: DimensionTypeOrDed = 1,
+    is_causal: Optional[CausalArgTypeOrDed] = False,
+    scale: Optional[float] = None,
+    q_tile_shape: Optional[DimensionType] = None,
+    kv_tile_shape: Optional[DimensionType] = None,
+    torch_compile: bool = False,
+    return_lse: bool = False,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    na_tensor_checks(
+        query, key, value, must_match_head_dims=True, supports_gqa_mqa=True
+    )
+    na_dim = query.dim() - 3  # batch, heads, head_dim
+    kernel_size, stride, dilation, is_causal = check_all_args(
+        na_dim, kernel_size, stride, dilation, is_causal
+    )
+    check_args_against_input(
+        query,
+        kernel_size=kernel_size,
+        stride=stride,
+        dilation=dilation,
+        is_causal=is_causal,
+    )
+    batch_size, *qkv_shape_in, num_heads, head_dim = query.shape
+    num_heads_kv, head_dim_v = value.shape[-2:]
+    qkv_shape = check_input_size_arg(na_dim, qkv_shape_in)
+    scale = scale or query.shape[-1] ** -0.5
+    assert can_run_flex_attention(
+        query, key, value, torch_compile=torch_compile, raise_error=True
+    )
+    if (q_tile_shape is None) ^ (kv_tile_shape is None):
+        raise ValueError(
+            "Please specify both q_tile_shape and kv_tile_shape, or neither one. "
+            f"Got {q_tile_shape=}, {kv_tile_shape=}."
+        )
+    do_token_permute = q_tile_shape is not None and kv_tile_shape is not None
+    q_shape = kv_shape = qkv_shape
+    q_tile_size: Optional[int] = None
+    kv_tile_size: Optional[int] = None
+    if do_token_permute:
+        q_tile_shape, kv_tile_shape = check_flex_fna_forward_config(
+            input_tensor=query,
+            q_tile_shape=q_tile_shape,
+            kv_tile_shape=kv_tile_shape,
+        )
+        q_tile_size = math.prod(q_tile_shape)
+        kv_tile_size = math.prod(kv_tile_shape)
+        assert q_tile_shape is not None
+        assert kv_tile_shape is not None
+        query_perm, _, q_shape = token_permute_operation(
+            query,
+            tile_shape=q_tile_shape,
+            dilation=dilation,
+            flip_tiled_dims=False,
+        )
+        key_perm, _, k_shape = token_permute_operation(
+            key, tile_shape=kv_tile_shape, dilation=dilation, flip_tiled_dims=False
+        )
+        value_perm, _, v_shape = token_permute_operation(
+            value, tile_shape=kv_tile_shape, dilation=dilation, flip_tiled_dims=False
+        )
+        assert k_shape == v_shape
+        kv_shape = k_shape
+        # Token permute already flattens to 1-D
+        # Flex uses heads first layout
+        query_ = query_perm.transpose(1, 2)
+        key_ = key_perm.transpose(1, 2)
+        value_ = value_perm.transpose(1, 2)
+    else:
+        seqlen = math.prod(qkv_shape)
+        # Flex uses heads first layout
+        query_ = query.reshape(batch_size, seqlen, num_heads, head_dim).transpose(1, 2)
+        key_ = key.reshape(batch_size, seqlen, num_heads_kv, head_dim).transpose(1, 2)
+        value_ = value.reshape(batch_size, seqlen, num_heads_kv, head_dim_v).transpose(
+            1, 2
+        )
+    na_block_mask = get_na_flex_mask(
+        device=query.device.type,
+        na_dim=na_dim,
+        qkv_shape=qkv_shape,
+        kernel_size=kernel_size,
+        stride=stride,
+        dilation=dilation,
+        is_causal=is_causal,
+        q_tile_shape=q_tile_shape,
+        kv_tile_shape=kv_tile_shape,
+        q_shape=q_shape,
+        kv_shape=kv_shape,
+        torch_compile=torch_compile,
+    )
+    out_, lse_ = run_flex_attn(
+        query_,
+        key_,
+        value_,
+        na_block_mask,
+        scale,
+        torch_compile=torch_compile,
+        q_tile_size=q_tile_size,
+        kv_tile_size=kv_tile_size,
+    )
+    if do_token_permute:
+        out = out_.transpose(1, 2)
+        lse = lse_.transpose(1, 2).unsqueeze(-1)
+        assert q_tile_shape is not None
+        assert kv_tile_shape is not None
+        out = token_unpermute_operation(
+            out,
+            token_layout_shape=qkv_shape,
+            tile_shape=q_tile_shape,
+            dilation=dilation,
+            flip_tiled_dims=False,
+        )
+        lse = token_unpermute_operation(
+            lse,
+            token_layout_shape=qkv_shape,
+            tile_shape=q_tile_shape,
+            dilation=dilation,
+            flip_tiled_dims=False,
+        ).squeeze(-1)
+    else:
+        out = out_.transpose(1, 2).reshape(
+            batch_size, *qkv_shape, num_heads, head_dim_v
+        )
+        lse = lse_.transpose(1, 2).reshape(batch_size, *qkv_shape, num_heads)
+    if return_lse:
+        return out, lse
+    return out
+def na1d_flex(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    kernel_size: Dimension1DTypeOrDed,
+    stride: Dimension1DTypeOrDed = 1,
+    dilation: Dimension1DTypeOrDed = 1,
+    is_causal: Optional[CausalArg1DTypeOrDed] = False,
+    scale: Optional[float] = None,
+    q_tile_shape: Optional[Dimension1DType] = None,
+    kv_tile_shape: Optional[Dimension1DType] = None,
+    torch_compile: bool = False,
+    return_lse: bool = False,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    return flex_fna_generic(
+        query=query,
+        key=key,
+        value=value,
+        kernel_size=kernel_size,
+        stride=stride,
+        dilation=dilation,
+        is_causal=is_causal,
+        scale=scale,
+        q_tile_shape=q_tile_shape,
+        kv_tile_shape=kv_tile_shape,
+        torch_compile=torch_compile,
+        return_lse=return_lse,
+    )
+def na2d_flex(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    kernel_size: Dimension2DTypeOrDed,
+    stride: Dimension2DTypeOrDed = 1,
+    dilation: Dimension2DTypeOrDed = 1,
+    is_causal: Optional[CausalArg2DTypeOrDed] = False,
+    scale: Optional[float] = None,
+    q_tile_shape: Optional[Dimension2DType] = None,
+    kv_tile_shape: Optional[Dimension2DType] = None,
+    torch_compile: bool = False,
+    return_lse: bool = False,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    return flex_fna_generic(
+        query=query,
+        key=key,
+        value=value,
+        kernel_size=kernel_size,
+        stride=stride,
+        dilation=dilation,
+        is_causal=is_causal,
+        scale=scale,
+        q_tile_shape=q_tile_shape,
+        kv_tile_shape=kv_tile_shape,
+        torch_compile=torch_compile,
+        return_lse=return_lse,
+    )
+def na3d_flex(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    kernel_size: Dimension3DTypeOrDed,
+    stride: Dimension3DTypeOrDed = 1,
+    dilation: Dimension3DTypeOrDed = 1,
+    is_causal: Optional[CausalArg3DTypeOrDed] = False,
+    scale: Optional[float] = None,
+    q_tile_shape: Optional[Dimension3DType] = None,
+    kv_tile_shape: Optional[Dimension3DType] = None,
+    torch_compile: bool = False,
+    return_lse: bool = False,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    return flex_fna_generic(
+        query=query,
+        key=key,
+        value=value,
+        kernel_size=kernel_size,
+        stride=stride,
+        dilation=dilation,
+        is_causal=is_causal,
+        scale=scale,
+        q_tile_shape=q_tile_shape,
+        kv_tile_shape=kv_tile_shape,
+        torch_compile=torch_compile,
+        return_lse=return_lse,
+    )

build/torch210-cu126-aarch64-linux/backends/fmha.py ADDED Viewed

	@@ -0,0 +1,283 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+import functools
+from typing import Optional, Tuple, Union
+import torch
+from torch import Tensor
+from torch.amp import custom_bwd, custom_fwd
+from torch.autograd import Function
+amp_fwd = functools.partial(custom_fwd, device_type="cuda")
+amp_bwd = functools.partial(custom_bwd, device_type="cuda")
+from .._libnatten import fmha_backward, fmha_forward
+from ..backends.configs.checks import can_run_cutlass_fmha
+from ..backends.configs.cutlass import (
+    check_cutlass_fmha_backward_config,
+    check_cutlass_fmha_forward_config,
+)
+from .._types import (
+    CutlassFmhaBackwardConfigType,
+    CutlassFmhaForwardConfigType,
+    NoneType,
+)
+from ..utils import log
+from ..utils.checks import fmha_tensor_checks, varlen_tensor_checks
+logger = log.get_logger(__name__)
+class CutlassFmhaAutogradFn(Function):
+    @staticmethod
+    @amp_fwd
+    def forward(
+        ctx,
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        is_causal: bool,
+        scale: float,
+        forward_config: CutlassFmhaForwardConfigType,
+        backward_config: CutlassFmhaBackwardConfigType,
+        backward_kv_splits: Optional[int],
+        backward_use_pt_reduction: bool,
+        cumulative_seqlen_Q: Optional[Tensor],
+        cumulative_seqlen_KV: Optional[Tensor],
+        max_seqlen_Q: int,
+        max_seqlen_KV: int,
+    ) -> Tuple[Tensor, Tensor]:
+        query = query.contiguous()
+        key = key.contiguous()
+        value = value.contiguous()
+        assert query.dim() == value.dim() == 4
+        assert query.shape[0] == value.shape[0]
+        assert query.shape[2] == value.shape[2]
+        q_tile_size, kv_tile_size = forward_config
+        output, logsumexp = fmha_forward(
+            query,
+            key,
+            value,
+            is_causal,
+            scale,
+            q_tile_size,
+            kv_tile_size,
+            cumulative_seqlen_Q,
+            cumulative_seqlen_KV,
+            max_seqlen_Q,
+            max_seqlen_KV,
+        )
+        ctx.save_for_backward(
+            query,
+            key,
+            value,
+            logsumexp,
+            output,
+            cumulative_seqlen_Q,
+            cumulative_seqlen_KV,
+        )
+        ctx.scale = scale
+        ctx.is_causal = is_causal
+        ctx.max_seqlen_Q = max_seqlen_Q
+        ctx.max_seqlen_KV = max_seqlen_KV
+        ctx.backward_config = backward_config
+        ctx.backward_kv_splits = backward_kv_splits
+        ctx.backward_use_pt_reduction = backward_use_pt_reduction
+        # Always record determinism behavior during forward pass (forward pass itself is
+        # deterministic anyway).
+        # Determinism could be limited to part of the program, which means during forward pass
+        # it'll be true, but on .backward() call, if it's been turned off, it will stay off when we
+        # get to this operation's backward call.
+        ctx.deterministic = torch.are_deterministic_algorithms_enabled()
+        return output, logsumexp
+    @staticmethod
+    @amp_bwd
+    def backward(ctx, grad_out: Tensor, grad_lse: Tensor) -> Tuple[
+        Tensor,
+        Tensor,
+        Tensor,
+        NoneType,
+        NoneType,
+        NoneType,
+        NoneType,
+        # kv_splits, use_pt_reduction
+        NoneType,
+        NoneType,
+        # varlen
+        NoneType,
+        NoneType,
+        NoneType,
+        NoneType,
+    ]:
+        (
+            query,
+            key,
+            value,
+            logsumexp,
+            output,
+            cumulative_seqlen_Q,
+            cumulative_seqlen_KV,
+        ) = ctx.saved_tensors
+        d_output = grad_out.contiguous()
+        q_tile_size, k_tile_size = ctx.backward_config
+        d_query, d_key, d_value = fmha_backward(
+            query,
+            key,
+            value,
+            output,
+            d_output,
+            logsumexp,
+            ctx.is_causal,
+            ctx.scale,
+            q_tile_size,
+            k_tile_size,
+            ctx.backward_kv_splits,
+            ctx.backward_use_pt_reduction,
+            cumulative_seqlen_Q,
+            cumulative_seqlen_KV,
+            ctx.max_seqlen_Q,
+            ctx.max_seqlen_KV,
+            ctx.deterministic,
+        )
+        return (
+            d_query,
+            d_key,
+            d_value,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+def cutlass_fmha(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    is_causal: bool = False,
+    scale: Optional[float] = None,
+    q_tile_size: Optional[int] = None,
+    kv_tile_size: Optional[int] = None,
+    backward_q_tile_size: Optional[int] = None,
+    backward_kv_tile_size: Optional[int] = None,
+    backward_kv_splits: Optional[int] = None,
+    backward_use_pt_reduction: bool = False,
+    return_lse: bool = False,
+    # varlen parameters
+    cumulative_seqlen_Q: Optional[Tensor] = None,
+    cumulative_seqlen_KV: Optional[Tensor] = None,
+    max_seqlen_Q: int = 0,
+    max_seqlen_KV: int = 0,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    fmha_tensor_checks(
+        query,
+        key,
+        value,
+        must_match_head_dims=False,
+        supports_gqa_mqa=True,
+        backend_name="CUTLASS FMHA",
+    )
+    (
+        cumulative_seqlen_Q,
+        cumulative_seqlen_KV,
+        max_seqlen_Q,
+        max_seqlen_KV,
+    ) = varlen_tensor_checks(
+        query=query,
+        key=key,
+        value=value,
+        cumulative_seqlen_Q=cumulative_seqlen_Q,
+        cumulative_seqlen_KV=cumulative_seqlen_KV,
+        max_seqlen_Q=max_seqlen_Q,
+        max_seqlen_KV=max_seqlen_KV,
+    )
+    is_varlen = cumulative_seqlen_Q is not None
+    assert can_run_cutlass_fmha(
+        query, key, value, is_causal=is_causal, is_varlen=is_varlen, raise_error=True
+    )
+    forward_config = check_cutlass_fmha_forward_config(
+        input_tensor=query if value.shape[-1] <= query.shape[-1] else value,
+        q_tile_size=q_tile_size,
+        kv_tile_size=kv_tile_size,
+    )
+    backward_config = check_cutlass_fmha_backward_config(
+        input_tensor=key if value.shape[-1] <= key.shape[-1] else value,
+        q_tile_size=backward_q_tile_size,
+        kv_tile_size=backward_kv_tile_size,
+    )
+    scale = scale or query.shape[-1] ** -0.5
+    # GQA/MQA is not supported by the kernel; only allowed via graph transform
+    is_gqa = query.shape[-2] != key.shape[-2]
+    if is_gqa:
+        heads = query.shape[-2]
+        heads_kv = key.shape[-2]
+        assert key.shape[-2] == value.shape[-2]
+        assert heads >= heads_kv
+        assert heads % heads_kv == 0
+        h_k = heads // heads_kv
+        key = torch.repeat_interleave(key, repeats=h_k, dim=-2, output_size=heads)
+        value = torch.repeat_interleave(value, repeats=h_k, dim=-2, output_size=heads)
+    output, lse = CutlassFmhaAutogradFn.apply(
+        query,
+        key,
+        value,
+        is_causal,
+        scale,
+        forward_config,
+        backward_config,
+        backward_kv_splits,
+        backward_use_pt_reduction,
+        cumulative_seqlen_Q,
+        cumulative_seqlen_KV,
+        max_seqlen_Q,
+        max_seqlen_KV,
+    )
+    if return_lse:
+        return output, lse
+    return output

build/torch210-cu126-aarch64-linux/backends/fna.py ADDED Viewed

	@@ -0,0 +1,417 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+import functools
+from typing import Optional, Tuple, Union
+import torch
+from torch import Tensor
+from torch.amp import custom_bwd, custom_fwd
+from torch.autograd import Function
+amp_fwd = functools.partial(custom_fwd, device_type="cuda")
+amp_bwd = functools.partial(custom_bwd, device_type="cuda")
+from .._libnatten import (
+    na1d_backward,
+    na1d_forward,
+    na2d_backward,
+    na2d_forward,
+    na3d_backward,
+    na3d_forward,
+)
+from ..backends.configs.checks import can_run_cutlass_fna
+from ..backends.configs.cutlass import (
+    check_cutlass_fna_backward_config,
+    check_cutlass_fna_forward_config,
+)
+from .._types import (
+    CausalArg1DTypeOrDed,
+    CausalArg2DTypeOrDed,
+    CausalArg3DTypeOrDed,
+    CausalArgType,
+    CausalArgTypeOrDed,
+    CutlassFnaBackwardConfigType,
+    CutlassFnaForwardConfigType,
+    Dimension1DType,
+    Dimension1DTypeOrDed,
+    Dimension2DType,
+    Dimension2DTypeOrDed,
+    Dimension3DType,
+    Dimension3DTypeOrDed,
+    DimensionType,
+    DimensionTypeOrDed,
+    NoneType,
+)
+from ..utils import log
+from ..utils.checks import (
+    check_all_args,
+    check_args_against_input,
+    na_tensor_checks,
+)
+logger = log.get_logger(__name__)
+def make_cutlass_fna_autograd_fn(na_dim):
+    assert na_dim in [1, 2, 3]
+    FORWARD_OPS = {
+        1: na1d_forward,
+        2: na2d_forward,
+        3: na3d_forward,
+    }
+    BACKWARD_OPS = {
+        1: na1d_backward,
+        2: na2d_backward,
+        3: na3d_backward,
+    }
+    class CutlassFnaGenericAutogradFn(Function):
+        @staticmethod
+        @amp_fwd
+        def forward(
+            ctx,
+            query: Tensor,
+            key: Tensor,
+            value: Tensor,
+            kernel_size: DimensionType,
+            stride: DimensionType,
+            dilation: DimensionType,
+            is_causal: CausalArgType,
+            scale: float,
+            forward_config: CutlassFnaForwardConfigType,
+            backward_config: CutlassFnaBackwardConfigType,
+            backward_kv_splits: Optional[DimensionType],
+            backward_use_pt_reduction: bool,
+        ) -> Tuple[Tensor, Tensor]:
+            kernel_size, stride, dilation, is_causal = check_all_args(
+                na_dim, kernel_size, stride, dilation, is_causal
+            )
+            query = query.contiguous()
+            key = key.contiguous()
+            value = value.contiguous()
+            assert query.dim() == value.dim() == 3 + na_dim
+            assert query.shape[0] == value.shape[0]
+            assert query.shape[-2] == value.shape[-2]
+            q_tile_shape, kv_tile_shape = forward_config
+            output, logsumexp = FORWARD_OPS[na_dim](
+                query,
+                key,
+                value,
+                kernel_size,
+                stride,
+                dilation,
+                is_causal,
+                scale,
+                q_tile_shape,
+                kv_tile_shape,
+            )
+            ctx.save_for_backward(query, key, value, logsumexp, output)
+            ctx.kernel_size = kernel_size
+            ctx.stride = stride
+            ctx.dilation = dilation
+            ctx.is_causal = is_causal
+            ctx.scale = scale
+            ctx.backward_config = backward_config
+            ctx.backward_kv_splits = backward_kv_splits
+            ctx.backward_use_pt_reduction = backward_use_pt_reduction
+            # Always record determinism behavior during forward pass (forward pass itself is
+            # deterministic anyway).
+            # Determinism could be limited to part of the program, which means during forward pass
+            # it'll be true, but on .backward() call, if it's been turned off, it will stay off when we
+            # get to this operation's backward call.
+            ctx.deterministic = torch.are_deterministic_algorithms_enabled()
+            return output, logsumexp
+        @staticmethod
+        @amp_bwd
+        def backward(ctx, grad_out: Tensor, grad_lse: Tensor) -> Tuple[
+            Tensor,
+            Tensor,
+            Tensor,
+            NoneType,
+            NoneType,
+            NoneType,
+            NoneType,
+            NoneType,
+            NoneType,
+            NoneType,
+            NoneType,
+            NoneType,
+        ]:
+            query, key, value, logsumexp, output = ctx.saved_tensors
+            d_output = grad_out.contiguous()
+            q_tile_shape, k_tile_shape = ctx.backward_config
+            d_query, d_key, d_value = BACKWARD_OPS[na_dim](
+                query,
+                key,
+                value,
+                output,
+                d_output,
+                logsumexp,
+                ctx.kernel_size,
+                ctx.stride,
+                ctx.dilation,
+                ctx.is_causal,
+                ctx.scale,
+                q_tile_shape,
+                k_tile_shape,
+                ctx.backward_kv_splits,
+                ctx.backward_use_pt_reduction,
+                ctx.deterministic,
+            )
+            return (
+                d_query,
+                d_key,
+                d_value,
+                None,
+                None,
+                None,
+                None,
+                None,
+                None,
+                None,
+                None,
+                None,
+            )
+    return CutlassFnaGenericAutogradFn
+CutlassFna1DAutogradFn = make_cutlass_fna_autograd_fn(1)
+CutlassFna2DAutogradFn = make_cutlass_fna_autograd_fn(2)
+CutlassFna3DAutogradFn = make_cutlass_fna_autograd_fn(3)
+CutlassFNAAutogradFns = {
+    1: CutlassFna1DAutogradFn,
+    2: CutlassFna2DAutogradFn,
+    3: CutlassFna3DAutogradFn,
+}
+def cutlass_fna_generic(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    kernel_size: DimensionTypeOrDed,
+    stride: DimensionTypeOrDed = 1,
+    dilation: DimensionTypeOrDed = 1,
+    is_causal: Optional[CausalArgTypeOrDed] = False,
+    scale: Optional[float] = None,
+    q_tile_shape: Optional[DimensionType] = None,
+    kv_tile_shape: Optional[DimensionType] = None,
+    backward_q_tile_shape: Optional[DimensionType] = None,
+    backward_kv_tile_shape: Optional[DimensionType] = None,
+    backward_kv_splits: Optional[DimensionType] = None,
+    backward_use_pt_reduction: bool = False,
+    return_lse: bool = False,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    na_tensor_checks(
+        query, key, value, must_match_head_dims=False, supports_gqa_mqa=True
+    )
+    assert can_run_cutlass_fna(query, key, value, raise_error=True)
+    na_dim = query.dim() - 3  # batch, heads, head_dim
+    assert na_dim in [1, 2, 3]
+    kernel_size, stride, dilation, is_causal = check_all_args(
+        na_dim, kernel_size, stride, dilation, is_causal
+    )
+    check_args_against_input(
+        query,
+        kernel_size=kernel_size,
+        stride=stride,
+        dilation=dilation,
+        is_causal=is_causal,
+    )
+    forward_config = check_cutlass_fna_forward_config(
+        input_tensor=query if value.shape[-1] <= query.shape[-1] else value,
+        dilation=dilation,
+        q_tile_shape=q_tile_shape,
+        kv_tile_shape=kv_tile_shape,
+    )
+    backward_config = check_cutlass_fna_backward_config(
+        input_tensor=key if value.shape[-1] <= key.shape[-1] else value,
+        q_tile_shape=backward_q_tile_shape,
+        kv_tile_shape=backward_kv_tile_shape,
+    )
+    scale = scale or query.shape[-1] ** -0.5
+    # GQA/MQA is not supported by the kernel; only allowed via graph transform
+    is_gqa = query.shape[-2] != key.shape[-2]
+    if is_gqa:
+        heads = query.shape[-2]
+        heads_kv = key.shape[-2]
+        assert key.shape[-2] == value.shape[-2]
+        assert heads >= heads_kv
+        assert heads % heads_kv == 0
+        h_k = heads // heads_kv
+        key = torch.repeat_interleave(key, repeats=h_k, dim=-2, output_size=heads)
+        value = torch.repeat_interleave(value, repeats=h_k, dim=-2, output_size=heads)
+    output, lse = CutlassFNAAutogradFns[na_dim].apply(
+        query,
+        key,
+        value,
+        kernel_size,
+        stride,
+        dilation,
+        is_causal,
+        scale,
+        forward_config,
+        backward_config,
+        backward_kv_splits,
+        backward_use_pt_reduction,
+    )
+    if return_lse:
+        return output, lse
+    return output
+def na1d_cutlass_fna(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    kernel_size: Dimension1DTypeOrDed,
+    stride: Dimension1DTypeOrDed = 1,
+    dilation: Dimension1DTypeOrDed = 1,
+    is_causal: Optional[CausalArg1DTypeOrDed] = False,
+    scale: Optional[float] = None,
+    q_tile_shape: Optional[Dimension1DType] = None,
+    kv_tile_shape: Optional[Dimension1DType] = None,
+    backward_q_tile_shape: Optional[Dimension1DType] = None,
+    backward_kv_tile_shape: Optional[Dimension1DType] = None,
+    backward_kv_splits: Optional[Dimension1DType] = None,
+    backward_use_pt_reduction: bool = False,
+    return_lse: bool = False,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    return cutlass_fna_generic(
+        query=query,
+        key=key,
+        value=value,
+        kernel_size=kernel_size,
+        stride=stride,
+        dilation=dilation,
+        is_causal=is_causal,
+        scale=scale,
+        q_tile_shape=q_tile_shape,
+        kv_tile_shape=kv_tile_shape,
+        backward_q_tile_shape=backward_q_tile_shape,
+        backward_kv_tile_shape=backward_kv_tile_shape,
+        backward_kv_splits=backward_kv_splits,
+        backward_use_pt_reduction=backward_use_pt_reduction,
+        return_lse=return_lse,
+    )
+def na2d_cutlass_fna(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    kernel_size: Dimension2DTypeOrDed,
+    stride: Dimension2DTypeOrDed = 1,
+    dilation: Dimension2DTypeOrDed = 1,
+    is_causal: Optional[CausalArg2DTypeOrDed] = False,
+    scale: Optional[float] = None,
+    q_tile_shape: Optional[Dimension2DType] = None,
+    kv_tile_shape: Optional[Dimension2DType] = None,
+    backward_q_tile_shape: Optional[Dimension2DType] = None,
+    backward_kv_tile_shape: Optional[Dimension2DType] = None,
+    backward_kv_splits: Optional[Dimension2DType] = None,
+    backward_use_pt_reduction: bool = False,
+    return_lse: bool = False,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    return cutlass_fna_generic(
+        query=query,
+        key=key,
+        value=value,
+        kernel_size=kernel_size,
+        stride=stride,
+        dilation=dilation,
+        is_causal=is_causal,
+        scale=scale,
+        q_tile_shape=q_tile_shape,
+        kv_tile_shape=kv_tile_shape,
+        backward_q_tile_shape=backward_q_tile_shape,
+        backward_kv_tile_shape=backward_kv_tile_shape,
+        backward_kv_splits=backward_kv_splits,
+        backward_use_pt_reduction=backward_use_pt_reduction,
+        return_lse=return_lse,
+    )
+def na3d_cutlass_fna(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    kernel_size: Dimension3DTypeOrDed,
+    stride: Dimension3DTypeOrDed = 1,
+    dilation: Dimension3DTypeOrDed = 1,
+    is_causal: Optional[CausalArg3DTypeOrDed] = False,
+    scale: Optional[float] = None,
+    q_tile_shape: Optional[Dimension3DType] = None,
+    kv_tile_shape: Optional[Dimension3DType] = None,
+    backward_q_tile_shape: Optional[Dimension3DType] = None,
+    backward_kv_tile_shape: Optional[Dimension3DType] = None,
+    backward_kv_splits: Optional[Dimension3DType] = None,
+    backward_use_pt_reduction: bool = False,
+    return_lse: bool = False,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    return cutlass_fna_generic(
+        query=query,
+        key=key,
+        value=value,
+        kernel_size=kernel_size,
+        stride=stride,
+        dilation=dilation,
+        is_causal=is_causal,
+        scale=scale,
+        q_tile_shape=q_tile_shape,
+        kv_tile_shape=kv_tile_shape,
+        backward_q_tile_shape=backward_q_tile_shape,
+        backward_kv_tile_shape=backward_kv_tile_shape,
+        backward_kv_splits=backward_kv_splits,
+        backward_use_pt_reduction=backward_use_pt_reduction,
+        return_lse=return_lse,
+    )

build/torch210-cu126-aarch64-linux/backends/hopper_fmha.py ADDED Viewed

	@@ -0,0 +1,261 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+import functools
+from typing import Optional, Tuple, Union
+import torch
+from torch import Tensor
+from torch.amp import custom_bwd, custom_fwd
+from torch.autograd import Function
+amp_fwd = functools.partial(custom_fwd, device_type="cuda")
+amp_bwd = functools.partial(custom_bwd, device_type="cuda")
+from .._libnatten import hopper_fmha_backward, hopper_fmha_forward
+from ..backends.configs.checks import can_run_cutlass_hopper_fmha
+from ..backends.configs.cutlass_hopper import (
+    check_cutlass_hopper_fmha_backward_config,
+    check_cutlass_hopper_fmha_forward_config,
+)
+from .._types import (
+    CutlassHopperFmhaBackwardConfigType,
+    CutlassHopperFmhaForwardConfigType,
+    KernelSchedule,
+    NoneType,
+)
+from ..utils import log
+from ..utils.checks import fmha_tensor_checks, varlen_tensor_checks
+logger = log.get_logger(__name__)
+class CutlassHopperFmhaAutogradFn(Function):
+    @staticmethod
+    @amp_fwd
+    def forward(
+        ctx,
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        is_causal: bool,
+        scale: float,
+        forward_config: CutlassHopperFmhaForwardConfigType,
+        backward_config: CutlassHopperFmhaBackwardConfigType,
+        cumulative_seqlen_Q: Optional[Tensor],
+        cumulative_seqlen_KV: Optional[Tensor],
+        max_seqlen_Q: int,
+        max_seqlen_KV: int,
+    ) -> Tuple[Tensor, Tensor]:
+        query = query.contiguous()
+        key = key.contiguous()
+        value = value.contiguous()
+        (q_tile_size, kv_tile_size), kernel_schedule = forward_config
+        output, logsumexp = hopper_fmha_forward(
+            query,
+            key,
+            value,
+            is_causal,
+            scale,
+            q_tile_size,
+            kv_tile_size,
+            kernel_schedule.value,  # TODO: I don't like this -- write a map with checks?
+            cumulative_seqlen_Q,
+            cumulative_seqlen_KV,
+            max_seqlen_Q,
+            max_seqlen_KV,
+        )
+        ctx.save_for_backward(
+            query,
+            key,
+            value,
+            logsumexp,
+            output,
+            cumulative_seqlen_Q,
+            cumulative_seqlen_KV,
+        )
+        ctx.scale = scale
+        ctx.is_causal = is_causal
+        ctx.max_seqlen_Q = max_seqlen_Q
+        ctx.max_seqlen_KV = max_seqlen_KV
+        ctx.backward_config = backward_config
+        # Always record determinism behavior during forward pass (forward pass itself is
+        # deterministic anyway).
+        # Determinism could be limited to part of the program, which means during forward pass
+        # it'll be true, but on .backward() call, if it's been turned off, it will stay off when we
+        # get to this operation's backward call.
+        ctx.deterministic = torch.are_deterministic_algorithms_enabled()
+        return output, logsumexp
+    @staticmethod
+    @amp_bwd
+    def backward(ctx, grad_out: Tensor, grad_lse: Tensor) -> Tuple[
+        Tensor,
+        Tensor,
+        Tensor,
+        NoneType,
+        NoneType,
+        NoneType,
+        NoneType,
+        NoneType,
+        NoneType,
+        NoneType,
+        NoneType,
+    ]:
+        (
+            query,
+            key,
+            value,
+            logsumexp,
+            output,
+            cumulative_seqlen_Q,
+            cumulative_seqlen_KV,
+        ) = ctx.saved_tensors
+        d_output = grad_out.contiguous()  # noqa: F841
+        q_tile_size, k_tile_size = ctx.backward_config
+        if ctx.deterministic:
+            raise RuntimeError(
+                "Hopper FMHA backward pass does not have a deterministic mode, "
+                "but PyTorch's deterministic algorithms were enabled. To proceed, "
+                "you must either disable torch's deterministic mode, or choose a "
+                "different backend."
+            )
+        d_query, d_key, d_value = hopper_fmha_backward(
+            query,
+            key,
+            value,
+            output,
+            d_output,
+            logsumexp,
+            ctx.is_causal,
+            ctx.scale,
+            q_tile_size,
+            k_tile_size,
+            cumulative_seqlen_Q,
+            cumulative_seqlen_KV,
+            ctx.max_seqlen_Q,
+            ctx.max_seqlen_KV,
+        )
+        return d_query, d_key, d_value, None, None, None, None, None, None, None, None
+def cutlass_hopper_fmha(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    is_causal: bool = False,
+    scale: Optional[float] = None,
+    q_tile_size: Optional[int] = None,
+    kv_tile_size: Optional[int] = None,
+    kernel_schedule: Optional[KernelSchedule] = None,
+    backward_q_tile_size: Optional[int] = None,
+    backward_kv_tile_size: Optional[int] = None,
+    return_lse: bool = False,
+    # varlen parameters
+    cumulative_seqlen_Q: Optional[Tensor] = None,
+    cumulative_seqlen_KV: Optional[Tensor] = None,
+    max_seqlen_Q: int = 0,
+    max_seqlen_KV: int = 0,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    fmha_tensor_checks(
+        query,
+        key,
+        value,
+        must_match_head_dims=True,
+        supports_gqa_mqa=True,
+        backend_name="Hopper FMHA",
+    )
+    (
+        cumulative_seqlen_Q,
+        cumulative_seqlen_KV,
+        max_seqlen_Q,
+        max_seqlen_KV,
+    ) = varlen_tensor_checks(
+        query=query,
+        key=key,
+        value=value,
+        cumulative_seqlen_Q=cumulative_seqlen_Q,
+        cumulative_seqlen_KV=cumulative_seqlen_KV,
+        max_seqlen_Q=max_seqlen_Q,
+        max_seqlen_KV=max_seqlen_KV,
+    )
+    is_varlen = cumulative_seqlen_Q is not None
+    assert can_run_cutlass_hopper_fmha(
+        query, key, value, is_causal=is_causal, is_varlen=is_varlen, raise_error=True
+    )
+    forward_config = check_cutlass_hopper_fmha_forward_config(
+        input_tensor=query,
+        q_tile_size=q_tile_size,
+        kv_tile_size=kv_tile_size,
+        kernel_schedule=kernel_schedule,
+    )
+    backward_config = check_cutlass_hopper_fmha_backward_config(
+        input_tensor=query,
+        q_tile_size=backward_q_tile_size,
+        kv_tile_size=backward_kv_tile_size,
+    )
+    scale = scale or query.shape[-1] ** -0.5
+    # GQA/MQA is not supported by the kernel; only allowed via graph transform
+    is_gqa = query.shape[-2] != key.shape[-2]
+    if is_gqa:
+        heads = query.shape[-2]
+        heads_kv = key.shape[-2]
+        assert key.shape[-2] == value.shape[-2]
+        assert heads >= heads_kv
+        assert heads % heads_kv == 0
+        h_k = heads // heads_kv
+        key = torch.repeat_interleave(key, repeats=h_k, dim=-2, output_size=heads)
+        value = torch.repeat_interleave(value, repeats=h_k, dim=-2, output_size=heads)
+    output, lse = CutlassHopperFmhaAutogradFn.apply(
+        query,
+        key,
+        value,
+        is_causal,
+        scale,
+        forward_config,
+        backward_config,
+        cumulative_seqlen_Q,
+        cumulative_seqlen_KV,
+        max_seqlen_Q,
+        max_seqlen_KV,
+    )
+    if return_lse:
+        return output, lse
+    return output

build/torch210-cu126-aarch64-linux/backends/hopper_fna.py ADDED Viewed

	@@ -0,0 +1,512 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+import functools
+from typing import Optional, Tuple, Union
+import torch
+from torch import Tensor
+from torch.amp import custom_bwd, custom_fwd
+from torch.autograd import Function
+amp_fwd = functools.partial(custom_fwd, device_type="cuda")
+amp_bwd = functools.partial(custom_bwd, device_type="cuda")
+from .._libnatten import (
+    hopper_na1d_backward,
+    hopper_na1d_forward,
+    hopper_na2d_backward,
+    hopper_na2d_forward,
+    hopper_na3d_backward,
+    hopper_na3d_forward,
+)
+from ..backends.configs.checks import can_run_cutlass_hopper_fna
+from ..backends.configs.cutlass_hopper import (
+    check_cutlass_hopper_fna_backward_config,
+    check_cutlass_hopper_fna_forward_config,
+)
+from ..token_permute import token_permute_operation, token_unpermute_operation
+from .._types import (
+    CausalArg1DTypeOrDed,
+    CausalArg2DTypeOrDed,
+    CausalArg3DTypeOrDed,
+    CausalArgType,
+    CausalArgTypeOrDed,
+    CutlassHopperFnaBackwardConfigType,
+    CutlassHopperFnaForwardConfigType,
+    Dimension1DType,
+    Dimension1DTypeOrDed,
+    Dimension2DType,
+    Dimension2DTypeOrDed,
+    Dimension3DType,
+    Dimension3DTypeOrDed,
+    DimensionType,
+    DimensionTypeOrDed,
+    KernelSchedule,
+    NoneType,
+)
+from ..utils.checks import (
+    check_all_args,
+    check_args_against_input,
+    na_tensor_checks,
+)
+def make_cutlass_hopper_fna_autograd_fn(na_dim):
+    assert na_dim in [1, 2, 3]
+    FORWARD_OPS = {
+        1: hopper_na1d_forward,
+        2: hopper_na2d_forward,
+        3: hopper_na3d_forward,
+    }
+    BACKWARD_OPS = {
+        1: hopper_na1d_backward,
+        2: hopper_na2d_backward,
+        3: hopper_na3d_backward,
+    }
+    class CutlassHopperFnaGenericAutogradFn(Function):
+        @staticmethod
+        @amp_fwd
+        def forward(
+            ctx,
+            query: Tensor,
+            key: Tensor,
+            value: Tensor,
+            kernel_size: DimensionType,
+            stride: DimensionType,
+            dilation: DimensionType,
+            is_causal: CausalArgType,
+            scale: float,
+            forward_config: CutlassHopperFnaForwardConfigType,
+            backward_config: CutlassHopperFnaBackwardConfigType,
+        ) -> Tuple[Tensor, Tensor]:
+            kernel_size, stride, dilation, is_causal = check_all_args(
+                na_dim, kernel_size, stride, dilation, is_causal
+            )
+            (q_tile_shape, kv_tile_shape), kernel_schedule = forward_config
+            # Token permute begin
+            query_perm, qkv_shape, q_shape = token_permute_operation(
+                query,
+                q_tile_shape,
+                dilation=dilation,
+                flip_tiled_dims=True,
+            )
+            key_perm, _, k_shape = token_permute_operation(
+                key, kv_tile_shape, dilation=dilation, flip_tiled_dims=True
+            )
+            value_perm, _, v_shape = token_permute_operation(
+                value, kv_tile_shape, dilation=dilation, flip_tiled_dims=True
+            )
+            assert k_shape == v_shape
+            kv_shape = k_shape
+            # Token permute end
+            query_perm = query_perm.contiguous()
+            key_perm = key_perm.contiguous()
+            value_perm = value_perm.contiguous()
+            output_perm, logsumexp_perm = FORWARD_OPS[na_dim](
+                query_perm,
+                key_perm,
+                value_perm,
+                kernel_size,
+                stride,
+                dilation,
+                is_causal,
+                scale,
+                q_shape,
+                kv_shape,
+                qkv_shape,
+                q_tile_shape,
+                kv_tile_shape,
+                kernel_schedule.value,  # TODO: I don't like this -- write a map with checks?
+            )
+            # Token un-permute begin
+            output = token_unpermute_operation(
+                output_perm,
+                token_layout_shape=qkv_shape,
+                tile_shape=q_tile_shape,
+                dilation=dilation,
+                flip_tiled_dims=True,
+            )
+            logsumexp = token_unpermute_operation(
+                logsumexp_perm.unsqueeze(-1),
+                token_layout_shape=qkv_shape,
+                tile_shape=q_tile_shape,
+                dilation=dilation,
+                flip_tiled_dims=True,
+            ).squeeze(-1)
+            # Token un-permute end
+            ctx.save_for_backward(query, key, value, logsumexp, output)
+            ctx.kernel_size = kernel_size
+            ctx.stride = stride
+            ctx.dilation = dilation
+            ctx.is_causal = is_causal
+            ctx.scale = scale
+            ctx.backward_config = backward_config
+            # Always record determinism behavior during forward pass (forward pass itself is
+            # deterministic anyway).
+            # Determinism could be limited to part of the program, which means during forward pass
+            # it'll be true, but on .backward() call, if it's been turned off, it will stay off when we
+            # get to this operation's backward call.
+            ctx.deterministic = torch.are_deterministic_algorithms_enabled()
+            return output, logsumexp
+        @staticmethod
+        @amp_bwd
+        def backward(ctx, d_output: Tensor, d_lse: Tensor) -> Tuple[
+            Tensor,
+            Tensor,
+            Tensor,
+            NoneType,
+            NoneType,
+            NoneType,
+            NoneType,
+            NoneType,
+            NoneType,
+            NoneType,
+        ]:
+            query, key, value, logsumexp, output = ctx.saved_tensors
+            kernel_size, stride, dilation, is_causal, scale = (
+                ctx.kernel_size,
+                ctx.stride,
+                ctx.dilation,
+                ctx.is_causal,
+                ctx.scale,
+            )
+            q_tile_shape, kv_tile_shape = ctx.backward_config
+            if ctx.deterministic:
+                raise RuntimeError(
+                    "Hopper FNA backward pass does not have a deterministic mode, "
+                    "but PyTorch's deterministic algorithms were enabled. To proceed, "
+                    "you must either disable torch's deterministic mode, or choose a "
+                    "different backend."
+                )
+            # Token permute begin
+            query_perm, qkv_shape, q_shape = token_permute_operation(
+                query, tile_shape=q_tile_shape, dilation=dilation, flip_tiled_dims=True
+            )
+            output_perm, _, o_shape = token_permute_operation(
+                output, tile_shape=q_tile_shape, dilation=dilation, flip_tiled_dims=True
+            )
+            d_output_perm, _, d_o_shape = token_permute_operation(
+                d_output,
+                tile_shape=q_tile_shape,
+                dilation=dilation,
+                flip_tiled_dims=True,
+            )
+            logsumexp_perm, _, _ = token_permute_operation(
+                logsumexp.unsqueeze(-1),
+                tile_shape=q_tile_shape,
+                dilation=dilation,
+                flip_tiled_dims=True,
+            )
+            key_perm, _, k_shape = token_permute_operation(
+                key, tile_shape=kv_tile_shape, dilation=dilation, flip_tiled_dims=True
+            )
+            value_perm, _, v_shape = token_permute_operation(
+                value, tile_shape=kv_tile_shape, dilation=dilation, flip_tiled_dims=True
+            )
+            assert q_shape == o_shape == d_o_shape
+            assert k_shape == v_shape
+            kv_shape = k_shape
+            # Token permute end
+            query_perm = query_perm.contiguous()
+            key_perm = key_perm.contiguous()
+            value_perm = value_perm.contiguous()
+            output_perm = output_perm.contiguous()
+            d_output_perm = d_output_perm.contiguous()
+            logsumexp_perm = logsumexp_perm.squeeze(-1)
+            d_query_perm, d_key_perm, d_value_perm = BACKWARD_OPS[na_dim](
+                query_perm,
+                key_perm,
+                value_perm,
+                output_perm,
+                d_output_perm,
+                logsumexp_perm,
+                kernel_size,
+                stride,
+                dilation,
+                is_causal,
+                scale,
+                q_shape,
+                kv_shape,
+                qkv_shape,
+                q_tile_shape,
+                kv_tile_shape,
+            )
+            # Token un-permute begin
+            d_query = token_unpermute_operation(
+                d_query_perm,
+                token_layout_shape=qkv_shape,
+                tile_shape=q_tile_shape,
+                dilation=dilation,
+                flip_tiled_dims=True,
+            )
+            d_key = token_unpermute_operation(
+                d_key_perm,
+                token_layout_shape=qkv_shape,
+                tile_shape=kv_tile_shape,
+                dilation=dilation,
+                flip_tiled_dims=True,
+            )
+            d_value = token_unpermute_operation(
+                d_value_perm,
+                token_layout_shape=qkv_shape,
+                tile_shape=kv_tile_shape,
+                dilation=dilation,
+                flip_tiled_dims=True,
+            )
+            # Token un-permute end
+            assert d_query.shape == query.shape
+            assert d_key.shape == key.shape
+            assert d_value.shape == value.shape
+            return (
+                d_query,
+                d_key,
+                d_value,
+                None,
+                None,
+                None,
+                None,
+                None,
+                None,
+                None,
+            )
+    return CutlassHopperFnaGenericAutogradFn
+CutlassHopperFna1DAutogradFn = make_cutlass_hopper_fna_autograd_fn(1)
+CutlassHopperFna2DAutogradFn = make_cutlass_hopper_fna_autograd_fn(2)
+CutlassHopperFna3DAutogradFn = make_cutlass_hopper_fna_autograd_fn(3)
+CutlassHopperFNAAutogradFns = {
+    1: CutlassHopperFna1DAutogradFn,
+    2: CutlassHopperFna2DAutogradFn,
+    3: CutlassHopperFna3DAutogradFn,
+}
+def cutlass_hopper_fna_generic(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    kernel_size: DimensionTypeOrDed,
+    stride: DimensionTypeOrDed = 1,
+    dilation: DimensionTypeOrDed = 1,
+    is_causal: Optional[CausalArgTypeOrDed] = False,
+    scale: Optional[float] = None,
+    q_tile_shape: Optional[DimensionType] = None,
+    kv_tile_shape: Optional[DimensionType] = None,
+    backward_q_tile_shape: Optional[DimensionType] = None,
+    backward_kv_tile_shape: Optional[DimensionType] = None,
+    kernel_schedule: Optional[KernelSchedule] = None,
+    return_lse: bool = False,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    na_tensor_checks(
+        query, key, value, must_match_head_dims=True, supports_gqa_mqa=True
+    )
+    assert can_run_cutlass_hopper_fna(query, key, value, raise_error=True)
+    na_dim = query.dim() - 3  # batch, heads, head_dim
+    kernel_size, stride, dilation, is_causal = check_all_args(
+        na_dim, kernel_size, stride, dilation, is_causal
+    )
+    check_args_against_input(
+        query,
+        kernel_size=kernel_size,
+        stride=stride,
+        dilation=dilation,
+        is_causal=is_causal,
+    )
+    forward_config = check_cutlass_hopper_fna_forward_config(
+        input_tensor=query,
+        q_tile_shape=q_tile_shape,
+        kv_tile_shape=kv_tile_shape,
+        kernel_schedule=kernel_schedule,
+    )
+    backward_config = check_cutlass_hopper_fna_backward_config(
+        input_tensor=query,
+        q_tile_shape=backward_q_tile_shape,
+        kv_tile_shape=backward_kv_tile_shape,
+    )
+    scale = scale or query.shape[-1] ** -0.5
+    # GQA/MQA is not supported by the kernel; only allowed via graph transform
+    is_gqa = query.shape[-2] != key.shape[-2]
+    if is_gqa:
+        heads = query.shape[-2]
+        heads_kv = key.shape[-2]
+        assert key.shape[-2] == value.shape[-2]
+        assert heads >= heads_kv
+        assert heads % heads_kv == 0
+        h_k = heads // heads_kv
+        key = torch.repeat_interleave(key, repeats=h_k, dim=-2, output_size=heads)
+        value = torch.repeat_interleave(value, repeats=h_k, dim=-2, output_size=heads)
+    output, lse = CutlassHopperFNAAutogradFns[na_dim].apply(
+        query,
+        key,
+        value,
+        kernel_size,
+        stride,
+        dilation,
+        is_causal,
+        scale,
+        forward_config,
+        backward_config,
+    )
+    if return_lse:
+        return output, lse
+    return output
+def na1d_cutlass_hopper_fna(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    kernel_size: Dimension1DTypeOrDed,
+    stride: Dimension1DTypeOrDed = 1,
+    dilation: Dimension1DTypeOrDed = 1,
+    is_causal: Optional[CausalArg1DTypeOrDed] = False,
+    scale: Optional[float] = None,
+    q_tile_shape: Optional[Dimension1DType] = None,
+    kv_tile_shape: Optional[Dimension1DType] = None,
+    backward_q_tile_shape: Optional[Dimension1DType] = None,
+    backward_kv_tile_shape: Optional[Dimension1DType] = None,
+    kernel_schedule: Optional[KernelSchedule] = None,
+    return_lse: bool = False,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    return cutlass_hopper_fna_generic(
+        query=query,
+        key=key,
+        value=value,
+        kernel_size=kernel_size,
+        stride=stride,
+        dilation=dilation,
+        is_causal=is_causal,
+        scale=scale,
+        q_tile_shape=q_tile_shape,
+        kv_tile_shape=kv_tile_shape,
+        backward_q_tile_shape=backward_q_tile_shape,
+        backward_kv_tile_shape=backward_kv_tile_shape,
+        kernel_schedule=kernel_schedule,
+        return_lse=return_lse,
+    )
+def na2d_cutlass_hopper_fna(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    kernel_size: Dimension2DTypeOrDed,
+    stride: Dimension2DTypeOrDed = 1,
+    dilation: Dimension2DTypeOrDed = 1,
+    is_causal: Optional[CausalArg2DTypeOrDed] = False,
+    scale: Optional[float] = None,
+    q_tile_shape: Optional[Dimension2DType] = None,
+    kv_tile_shape: Optional[Dimension2DType] = None,
+    backward_q_tile_shape: Optional[Dimension2DType] = None,
+    backward_kv_tile_shape: Optional[Dimension2DType] = None,
+    kernel_schedule: Optional[KernelSchedule] = None,
+    return_lse: bool = False,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    return cutlass_hopper_fna_generic(
+        query=query,
+        key=key,
+        value=value,
+        kernel_size=kernel_size,
+        stride=stride,
+        dilation=dilation,
+        is_causal=is_causal,
+        scale=scale,
+        q_tile_shape=q_tile_shape,
+        kv_tile_shape=kv_tile_shape,
+        backward_q_tile_shape=backward_q_tile_shape,
+        backward_kv_tile_shape=backward_kv_tile_shape,
+        kernel_schedule=kernel_schedule,
+        return_lse=return_lse,
+    )
+def na3d_cutlass_hopper_fna(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    kernel_size: Dimension3DTypeOrDed,
+    stride: Dimension3DTypeOrDed = 1,
+    dilation: Dimension3DTypeOrDed = 1,
+    is_causal: Optional[CausalArg3DTypeOrDed] = False,
+    scale: Optional[float] = None,
+    q_tile_shape: Optional[Dimension3DType] = None,
+    kv_tile_shape: Optional[Dimension3DType] = None,
+    backward_q_tile_shape: Optional[Dimension3DType] = None,
+    backward_kv_tile_shape: Optional[Dimension3DType] = None,
+    kernel_schedule: Optional[KernelSchedule] = None,
+    return_lse: bool = False,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    return cutlass_hopper_fna_generic(
+        query=query,
+        key=key,
+        value=value,
+        kernel_size=kernel_size,
+        stride=stride,
+        dilation=dilation,
+        is_causal=is_causal,
+        scale=scale,
+        q_tile_shape=q_tile_shape,
+        kv_tile_shape=kv_tile_shape,
+        backward_q_tile_shape=backward_q_tile_shape,
+        backward_kv_tile_shape=backward_kv_tile_shape,
+        kernel_schedule=kernel_schedule,
+        return_lse=return_lse,
+    )

build/torch210-cu126-aarch64-linux/backends/reference.py ADDED Viewed

	@@ -0,0 +1,343 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+import functools
+from typing import Optional, Tuple, Union
+import torch
+from torch import Tensor
+from torch.amp import custom_bwd, custom_fwd
+from torch.autograd import Function
+amp_fwd = functools.partial(custom_fwd, device_type="cuda")
+amp_bwd = functools.partial(custom_bwd, device_type="cuda")
+from .._libnatten import (
+    reference_na1d_backward,
+    reference_na1d_forward,
+    reference_na2d_backward,
+    reference_na2d_forward,
+    reference_na3d_backward,
+    reference_na3d_forward,
+)
+from .._types import (
+    CausalArg1DTypeOrDed,
+    CausalArg2DTypeOrDed,
+    CausalArg3DTypeOrDed,
+    CausalArgType,
+    CausalArgTypeOrDed,
+    Dimension1DTypeOrDed,
+    Dimension2DTypeOrDed,
+    Dimension3DTypeOrDed,
+    DimensionType,
+    DimensionTypeOrDed,
+    NoneType,
+)
+from ..utils import log
+from ..utils.checks import (
+    additional_kv_tensor_checks,
+    check_all_args,
+    check_args_against_input,
+    na_tensor_checks,
+)
+logger = log.get_logger(__name__)
+def make_reference_fna_autograd_fn(na_dim):
+    assert na_dim in [1, 2, 3]
+    FORWARD_OPS = {
+        1: reference_na1d_forward,
+        2: reference_na2d_forward,
+        3: reference_na3d_forward,
+    }
+    BACKWARD_OPS = {
+        1: reference_na1d_backward,
+        2: reference_na2d_backward,
+        3: reference_na3d_backward,
+    }
+    class ReferenceFnaGenericAutogradFn(Function):
+        @staticmethod
+        @amp_fwd
+        def forward(
+            ctx,
+            query: Tensor,
+            key: Tensor,
+            value: Tensor,
+            kernel_size: DimensionType,
+            stride: DimensionType,
+            dilation: DimensionType,
+            is_causal: CausalArgType,
+            scale: float,
+            qkv_shape: DimensionType,
+            num_extra_kv: int,
+        ) -> Tuple[Tensor, Tensor]:
+            kernel_size, stride, dilation, is_causal = check_all_args(
+                na_dim, kernel_size, stride, dilation, is_causal
+            )
+            query = query.contiguous()
+            key = key.contiguous()
+            value = value.contiguous()
+            assert query.dim() == value.dim() == 4
+            assert query.shape[0] == value.shape[0]
+            output, logsumexp = FORWARD_OPS[na_dim](
+                query,
+                key,
+                value,
+                kernel_size,
+                stride,
+                dilation,
+                is_causal,
+                scale,
+                qkv_shape,
+                num_extra_kv,
+            )
+            ctx.save_for_backward(query, key, value, logsumexp, output)
+            ctx.kernel_size = kernel_size
+            ctx.stride = stride
+            ctx.dilation = dilation
+            ctx.is_causal = is_causal
+            ctx.scale = scale
+            ctx.qkv_shape = qkv_shape
+            ctx.num_extra_kv = num_extra_kv
+            return output, logsumexp
+        @staticmethod
+        @amp_bwd
+        def backward(ctx, grad_out: Tensor, grad_lse: Tensor) -> Tuple[
+            Tensor,
+            Tensor,
+            Tensor,
+            NoneType,
+            NoneType,
+            NoneType,
+            NoneType,
+            NoneType,
+            NoneType,
+            NoneType,
+        ]:
+            query, key, value, logsumexp, output = ctx.saved_tensors
+            d_output = grad_out.contiguous()
+            d_query, d_key, d_value = BACKWARD_OPS[na_dim](
+                query,
+                key,
+                value,
+                output,
+                d_output,
+                logsumexp,
+                ctx.kernel_size,
+                ctx.stride,
+                ctx.dilation,
+                ctx.is_causal,
+                ctx.scale,
+                ctx.qkv_shape,
+                ctx.num_extra_kv,
+            )
+            return d_query, d_key, d_value, None, None, None, None, None, None, None
+    return ReferenceFnaGenericAutogradFn
+ReferenceFna1DAutogradFn = make_reference_fna_autograd_fn(1)
+ReferenceFna2DAutogradFn = make_reference_fna_autograd_fn(2)
+ReferenceFna3DAutogradFn = make_reference_fna_autograd_fn(3)
+ReferenceFnaAutogradFns = {
+    1: ReferenceFna1DAutogradFn,
+    2: ReferenceFna2DAutogradFn,
+    3: ReferenceFna3DAutogradFn,
+}
+def reference_fna_generic(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    kernel_size: DimensionTypeOrDed,
+    stride: DimensionTypeOrDed = 1,
+    dilation: DimensionTypeOrDed = 1,
+    is_causal: Optional[CausalArgTypeOrDed] = False,
+    scale: Optional[float] = None,
+    additional_keys: Optional[Tensor] = None,
+    additional_values: Optional[Tensor] = None,
+    return_lse: bool = False,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    na_tensor_checks(
+        query, key, value, must_match_head_dims=False, supports_gqa_mqa=True
+    )
+    additional_kv_tensor_checks(
+        query,
+        key,
+        value,
+        additional_keys,
+        additional_values,
+        must_match_head_dims=False,
+        supports_gqa_mqa=True,
+    )
+    na_dim = query.dim() - 3  # batch, heads, head_dim
+    assert na_dim in [1, 2, 3]
+    kernel_size, stride, dilation, is_causal = check_all_args(
+        na_dim, kernel_size, stride, dilation, is_causal
+    )
+    check_args_against_input(
+        query,
+        kernel_size=kernel_size,
+        stride=stride,
+        dilation=dilation,
+        is_causal=is_causal,
+    )
+    scale = scale or query.shape[-1] ** -0.5
+    qkv_shape = query.shape[1 : 1 + na_dim]
+    query = query.flatten(1, na_dim)
+    key = key.flatten(1, na_dim)
+    value = value.flatten(1, na_dim)
+    num_extra_kv = 0
+    if additional_keys is not None and additional_values is not None:
+        num_extra_kv = additional_keys.shape[1]
+        key = torch.cat([key, additional_keys], dim=1)
+        value = torch.cat([value, additional_values], dim=1)
+    output, lse = ReferenceFnaAutogradFns[na_dim].apply(
+        query,
+        key,
+        value,
+        kernel_size,
+        stride,
+        dilation,
+        is_causal,
+        scale,
+        qkv_shape,
+        num_extra_kv,
+    )
+    output = output.reshape(
+        query.shape[0], *qkv_shape, query.shape[-2], value.shape[-1]
+    )
+    lse = lse.reshape(query.shape[0], *qkv_shape, query.shape[-2])
+    if return_lse:
+        return output, lse
+    return output
+def na1d_reference(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    kernel_size: Dimension1DTypeOrDed,
+    stride: Dimension1DTypeOrDed = 1,
+    dilation: Dimension1DTypeOrDed = 1,
+    is_causal: Optional[CausalArg1DTypeOrDed] = False,
+    scale: Optional[float] = None,
+    additional_keys: Optional[Tensor] = None,
+    additional_values: Optional[Tensor] = None,
+    return_lse: bool = False,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    return reference_fna_generic(
+        query=query,
+        key=key,
+        value=value,
+        kernel_size=kernel_size,
+        stride=stride,
+        dilation=dilation,
+        is_causal=is_causal,
+        scale=scale,
+        additional_keys=additional_keys,
+        additional_values=additional_values,
+        return_lse=return_lse,
+    )
+def na2d_reference(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    kernel_size: Dimension2DTypeOrDed,
+    stride: Dimension2DTypeOrDed = 1,
+    dilation: Dimension2DTypeOrDed = 1,
+    is_causal: Optional[CausalArg2DTypeOrDed] = False,
+    scale: Optional[float] = None,
+    additional_keys: Optional[Tensor] = None,
+    additional_values: Optional[Tensor] = None,
+    return_lse: bool = False,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    return reference_fna_generic(
+        query=query,
+        key=key,
+        value=value,
+        kernel_size=kernel_size,
+        stride=stride,
+        dilation=dilation,
+        is_causal=is_causal,
+        scale=scale,
+        additional_keys=additional_keys,
+        additional_values=additional_values,
+        return_lse=return_lse,
+    )
+def na3d_reference(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    kernel_size: Dimension3DTypeOrDed,
+    stride: Dimension3DTypeOrDed = 1,
+    dilation: Dimension3DTypeOrDed = 1,
+    is_causal: Optional[CausalArg3DTypeOrDed] = False,
+    scale: Optional[float] = None,
+    additional_keys: Optional[Tensor] = None,
+    additional_values: Optional[Tensor] = None,
+    return_lse: bool = False,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    return reference_fna_generic(
+        query=query,
+        key=key,
+        value=value,
+        kernel_size=kernel_size,
+        stride=stride,
+        dilation=dilation,
+        is_causal=is_causal,
+        scale=scale,
+        additional_keys=additional_keys,
+        additional_values=additional_values,
+        return_lse=return_lse,
+    )

build/torch210-cu126-aarch64-linux/context.py ADDED Viewed

	@@ -0,0 +1,231 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+from enum import Enum
+import torch
+from .utils import log
+logger = log.get_logger(__name__)
+class MemoryUsagePreference(Enum):
+    Default = 0
+    Strict = 1
+    Unrestricted = 2
+class NattenContext:
+    is_deterministic_mode_enabled: bool = False
+    is_kv_parallelism_enabled: bool = True
+    training_memory_preference: MemoryUsagePreference = MemoryUsagePreference.Default
+    flex_compile_allowed: bool = False
+    flex_compile_backprop_allowed: bool = False
+    @staticmethod
+    def reset():
+        NattenContext.is_deterministic_mode_enabled = False
+        NattenContext.is_kv_parallelism_enabled = True
+        NattenContext.training_memory_preference = MemoryUsagePreference.Default
+        NattenContext.flex_compile_allowed = False
+        NattenContext.flex_compile_backprop_allowed = False
+def set_memory_usage_preference(pref: str = "default"):
+    """Sets memory usage preference for KV parallelism in `"cutlass-fna"` and `"cutlass-fmha"`
+    backends.
+    Args:
+        pref: Choices are `"default"`, `"strict"`, and `"unrestricted"`.
+    """
+    if pref == "default":
+        NattenContext.training_memory_preference = MemoryUsagePreference.Default
+    elif pref == "strict":
+        NattenContext.training_memory_preference = MemoryUsagePreference.Strict
+    elif pref == "unrestricted":
+        NattenContext.training_memory_preference = MemoryUsagePreference.Unrestricted
+    else:
+        raise ValueError(
+            "natten.set_memory_usage_preference allows only one of three settings: "
+            "`default`, `strict`, and `unrestricted`."
+        )
+def get_memory_usage_preference() -> MemoryUsagePreference:
+    return NattenContext.training_memory_preference
+def is_memory_usage_default() -> bool:
+    """Returns whether memory usage preference for KV parallelism in `"cutlass-fna"` and
+    `"cutlass-fmha"` backends is the default setting.
+    """
+    return get_memory_usage_preference() == MemoryUsagePreference.Default
+def is_memory_usage_strict() -> bool:
+    """Returns whether memory usage preference for KV parallelism in `"cutlass-fna"` and
+    `"cutlass-fmha"` backends is the *restricted* setting.
+    """
+    return get_memory_usage_preference() == MemoryUsagePreference.Strict
+def is_memory_usage_unrestricted() -> bool:
+    """Returns whether memory usage preference for KV parallelism in `"cutlass-fna"` and
+    `"cutlass-fmha"` backends is the *unrestricted* setting.
+    """
+    return get_memory_usage_preference() == MemoryUsagePreference.Unrestricted
+def use_deterministic_algorithms(mode: bool = True):
+    NattenContext.is_deterministic_mode_enabled = mode
+    if mode:
+        logger.warning(
+            "You're enabling NATTEN's deterministic mode. This mode does not "
+            "support auto-tuning, or training with positional biases. "
+            "For more information please refer to https://github.com/SHI-Labs/NATTEN/tree/main/docs"
+        )
+def are_deterministic_algorithms_enabled() -> bool:
+    return NattenContext.is_deterministic_mode_enabled
+def use_kv_parallelism_in_fused_na(mode: bool = True):
+    """Sets guards for using KV Parallelism in backpropagation in `"cutlass-fna"`/`"cutlass-fmha"`
+    backends.
+    Warning:
+        Disabling KV parallelism can significantly slow down training, particularly in
+        small-batch/head and large-token problems.
+    Args:
+        mode: If `True`, allows KV parallelism (default setting), and otherwise disables it.
+    """
+    if not mode:
+        NattenContext.is_kv_parallelism_enabled = False
+        return
+    if torch.are_deterministic_algorithms_enabled():
+        logger.warning(
+            "Attempted to enable KV parallelism in FNA, which is non-deterministic, "
+            "but PyTorch's deterministic flag has been enabled. Ignoring..."
+        )
+        return
+    if are_deterministic_algorithms_enabled():
+        raise RuntimeError(
+            "You enabled NATTEN's deterministic mode, but attempted to "
+            "enable KV parallelism, which results in non-determinism. "
+        )
+    NattenContext.is_kv_parallelism_enabled = True
+def is_kv_parallelism_in_fused_na_enabled() -> bool:
+    """Returns whether KV parallelism in `"cutlass-fna"` and `"cutlass-fmha"` backends is enabled."""
+    return NattenContext.is_kv_parallelism_enabled
+def is_flex_compile_allowed() -> bool:
+    """Returns whether compilation is allowed in `"flex-fna"` and `"flex-fmha"` backends."""
+    return NattenContext.flex_compile_allowed
+def is_flex_compile_backprop_allowed() -> bool:
+    """Returns whether compilation for backpropagation is allowed in `"flex-fna"` and `"flex-fmha"`
+    backends.
+    """
+    return NattenContext.flex_compile_backprop_allowed
+def allow_flex_compile(mode: bool = True, backprop: bool = False):
+    """Sets guards for Flex Attention + `torch.compile`.
+    Allows using our Flex FNA / Flex FMHA backends with `torch.compile`, meaning you can
+    pass `torch_compile=True` to the `na{1,2,3}d` or `attention` operation, along with
+    `backend="flex-fna"`/`backend="flex-fmha"`, and NATTEN will compile the block-sparse mask, as
+    well as the attention operation using `torch.compile` for you.
+    Warning:
+        We have been *unable to verify the correctness* of this setting under all of our use
+        cases. We are working on raising this issue with PyTorch directly, but until then we strongly
+        recommend exercising caution when using this feature.
+    Danger: backprop=True is strongly discouraged!
+        Allowing `torch.compile` for backpropagation (detected by checking
+        `tensor.requires_grad`) is guarded separately. We strongly recommend NOT using this setting, as
+        it can impact your training results.
+    Args:
+        mode: If `True`, enable compilation for forward pass, otherwise disable.
+        backprop: If `True`, assuming compilation for forward pass is allowed, enable compilation
+            for backward pass, otherwise disable.
+    """
+    if not mode:
+        NattenContext.flex_compile_allowed = False
+        NattenContext.flex_compile_backprop_allowed = False
+    if not NattenContext.flex_compile_allowed:
+        logger.warning(
+            "You are enabling Flex Attention compilation in NATTEN. "
+            "NATTEN does not allow this by default, because we cannot verify Flex's correctness in all "
+            "scenarios through NATTEN's tests. By choosing to override this, you acknowledge that your "
+            "results may be affected significantly. If this was not intended, please call "
+            "natten.disable_flex_compile()"
+            ""
+        )
+    NattenContext.flex_compile_allowed = True
+    if backprop:
+        if not NattenContext.flex_compile_backprop_allowed:
+            logger.warning(
+                "You are enabling using compiled Flex Attention to backpropagate. "
+                "NATTEN does not allow this by default, because we cannot verify Flex's correctness in all "
+                "scenarios through NATTEN's tests, and it is HIGHLY discouraged. By choosing to override "
+                "this, you acknowledge that your results may be heavily impacted significantly. "
+                "If this was not intended, please call "
+                "natten.disable_flex_compile_backprop()"
+                ""
+            )
+        NattenContext.flex_compile_backprop_allowed = True
+def allow_flex_compile_backprop(mode: bool = True):
+    """Sets guards for Flex Attention + `torch.compile` for backpropagation only.
+    Args:
+        mode: If `True`, enable compilation for backprop (assuming forward compilation is already
+            enabled), otherwise disable.
+    """
+    return allow_flex_compile(is_flex_compile_allowed(), mode)
+def disable_flex_compile():
+    """Disallow Flex Attention + `torch.compile` entirely."""
+    return allow_flex_compile(False)
+def disable_flex_compile_backprop():
+    """Disallow Flex Attention + `torch.compile` for backpropagation entirely."""
+    return allow_flex_compile(is_flex_compile_allowed(), False)

build/torch210-cu126-aarch64-linux/functional.py ADDED Viewed

	@@ -0,0 +1,1151 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+from typing import Dict, Optional, Tuple, Union
+import torch
+from torch import Tensor
+from .attn_merge import merge_attentions
+from .backends import (
+    choose_backend,
+    choose_fmha_backend,
+    cutlass_blackwell_fmha,
+    cutlass_blackwell_fna_generic,
+    cutlass_fmha,
+    cutlass_fna_generic,
+    cutlass_hopper_fmha,
+    cutlass_hopper_fna_generic,
+    flex_fmha,
+    flex_fna_generic,
+)
+from ._types import (
+    CausalArg1DTypeOrDed,
+    CausalArg2DTypeOrDed,
+    CausalArg3DTypeOrDed,
+    CausalArgTypeOrDed,
+    Dimension1DType,
+    Dimension1DTypeOrDed,
+    Dimension2DType,
+    Dimension2DTypeOrDed,
+    Dimension3DType,
+    Dimension3DTypeOrDed,
+    DimensionType,
+    DimensionTypeOrDed,
+    KernelSchedule,
+)
+from .utils import log
+from .utils.checks import (
+    additional_kv_tensor_checks,
+    check_all_args,
+    check_args_against_input,
+    check_kernel_schedule,
+    fmha_tensor_checks,
+    is_self_attention,
+    na_tensor_checks,
+    varlen_tensor_checks,
+)
+logger = log.get_logger(__name__)
+# Standard Attention
+def attention(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    is_causal: bool = False,
+    scale: Optional[float] = None,
+    # varlen parameters
+    seqlens_Q: Optional[Tensor] = None,
+    seqlens_KV: Optional[Tensor] = None,
+    cumulative_seqlen_Q: Optional[Tensor] = None,
+    cumulative_seqlen_KV: Optional[Tensor] = None,
+    max_seqlen_Q: Optional[int] = None,
+    max_seqlen_KV: Optional[int] = None,
+    # backend parameters
+    backend: Optional[str] = None,
+    q_tile_size: Optional[int] = None,
+    kv_tile_size: Optional[int] = None,
+    backward_q_tile_size: Optional[int] = None,
+    backward_kv_tile_size: Optional[int] = None,
+    backward_kv_splits: Optional[int] = None,
+    backward_use_pt_reduction: bool = False,
+    run_persistent_kernel: bool = True,
+    kernel_schedule: Optional[Union[str, KernelSchedule]] = None,
+    torch_compile: bool = False,
+    return_lse: bool = False,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    """Runs standard dot product attention.
+    This operation is used to implement neighborhood cross attention, in which we allow every
+    token to interact with some additional context (`additional_keys` and `additional_values`
+    tensors in [na1d][natten.na1d], [na2d][natten.na2d], and [na3d][natten.na3d]).
+    This operator is also used as a fast path for cases where neighborhood attention is equivalent
+    to self attention (not causal along any dims, and `kernel_size` is equal to the number of input
+    tokens).
+    This operation does not call into PyTorch's SDPA, and only runs one of the NATTEN backends
+    (`cutlass-fmha`, `hopper-fmha`, `blackwell-fmha`, `flex-fmha`). Reasons for that include being
+    able to control performance-related arguments, return logsumexp, and more.
+    For more information refer to [backends](backends.md).
+    Causal mask, and Variable length (varlen) Attention are also supported in some backends
+    (`cutlass-fmha` and `blackwell-fmha`).
+    Varlen Attention is only supported for the sequence-packed layout: QKV tensors have batch size
+    1, and tokens from different batches are concatenated without any padding along the sequence
+    dimension. Sequence lengths for different batches can be provided in two ways:
+        1. `seqlens_Q` and `seqlens_KV` (less efficient): only provide the sequence lengths as
+            integer tensors (must be on the same device as QKV), and NATTEN will compute cumulative
+            and maximum sequence lengths on each call.
+            This is **incompatible** with full-graph `torch.compile` since it requires a
+            synchronization.
+        2. `cumulative_seqlen_{Q,KV}` and `max_seqlen_{Q,KV}` (more efficient):
+            compute cumulative and maximum sequence lengths. `cumulative_seqlen_{Q,KV}` are integer
+            tensors on the same device as QKV containing the cumulative sum of `seqlens_{Q,KV}`,
+            with an additional `0` element in the beginning, therefore sized `batch+1`.
+            `max_seqlen_{Q,KV}` are integers (not Tensors) that represent the maximum sequence
+            lengths for Q and KV among all sequence batches.
+            You can use `natten.utils.varlen.generate_varlen_parameters` to generate these
+            parameters:
+                ```python3
+                from .utils.varlen import generate_varlen_parameters
+                (
+                    cumulative_seqlen_Q,
+                    cumulative_seqlen_KV,
+                    max_seqlen_Q,
+                    max_seqlen_KV,
+                ) = generate_varlen_parameters(q, k, v, seqlens_Q, seqlens_KV)
+                ```
+            As long as `generate_varlen_parameters` is called ahead of torch.compiling the model, it
+            is supported without any graph breaks.
+    GQA/MQA support (`heads != heads_kv`) is available. For now, `blackwell-fmha` and
+    `flex-fmha` support GQA/MQA natively, and `cutlass-fmha` and `hopper-fmha` support it with
+    explicit repeats (increases memory usage and runtime).
+    Parameters:
+        query (Tensor): 4-D query tensor, with the heads last layout
+            (`[batch, seqlen, heads, head_dim]`)
+        key (Tensor): 4-D key tensor, with the heads last layout
+            (`[batch, seqlen_kv, heads_kv, head_dim]`)
+        value (Tensor): 4-D value tensor, with the heads last layout
+            (`[batch, seqlen_kv, heads_kv, head_dim_v]`)
+        is_causal (bool): Toggle causal masking. Defaults to `False` (bi-directional).
+        scale (float): Attention scale. Defaults to `head_dim ** -0.5`.
+        seqlens_Q (Optional[Tensor]): (varlen) Optional 1-D tensor with size `batch`
+            indicating the number of query tokens in each batch. Must be passed together with
+            `seqlens_KV`.
+        seqlens_KV (Optional[Tensor]): (varlen) Optional 1-D tensor with size `batch`
+            indicating the number of key/value tokens in each batch. Must be passed together with
+            `seqlens_Q`.
+        cumulative_seqlen_Q (Optional[Tensor]): (varlen) Optional 1-D tensor with size `batch + 1`
+            indicating the cumulative sum of number of query tokens in each batch, with an
+            additional 0 element in the beginning. Must be passed together with
+            `cumulative_seqlen_KV` and `max_seqlen_{Q,KV}`.
+        cumulative_seqlen_KV (Optional[Tensor]): (varlen) Optional 1-D tensor with size `batch + 1`
+            indicating the cumulative sum of number of key/value tokens in each batch, with an
+            additional 0 element in the beginning. Must be passed together with
+            `cumulative_seqlen_Q` and `max_seqlen_{Q,KV}`.
+        max_seqlen_Q (Optional[int]): (varlen) Optional integer indicating the maximum query
+            sequence length in all batches. Must be passed together with `cumulative_seqlen_{Q,KV}`
+            and `max_seqlen_KV`.
+        max_seqlen_KV (Optional[int]): (varlen) Optional integer indicating the maximum key/value
+            sequence length in all batches. Must be passed together with `cumulative_seqlen_{Q,KV}`
+            and `max_seqlen_Q`.
+    Other Parameters:
+        backend (str): Backend implementation to run with. Choices are: `None` (pick the best
+            available one), `"cutlass-fmha"`, `"hopper-fmha"`, `"blackwell-fmha"`, `"flex-fmha"`.
+            Refer to [backends](backends.md) for more information.
+        q_tile_size (int): Tile size along query sequence length in the forward pass kernel.
+            You can use [profiler](profiler.md) to find valid choices for your use case.
+        kv_tile_size (int): Tile size along key/value sequence length in the forward pass kernel.
+            You can use [profiler](profiler.md) to find valid choices for your use case.
+        backward_q_tile_size (int): Tile size along query sequence length in the backward pass
+            kernel. This is ignored by `"flex-fmha"`.
+            You can use [profiler](profiler.md) to find valid choices for your use case.
+        backward_kv_tile_size (int): Tile size along key/value sequence length in the backward pass
+            kernel. This is ignored by `"flex-fmha"`.
+            You can use [profiler](profiler.md) to find valid choices for your use case.
+        backward_kv_splits (int): Number of key/value tiles allowed to work in parallel in the
+            backward pass kernel. This is only respected by the `"cutlass-fmha"` backend, only when
+            [KV parallelism](context.md#kv-parallelism-in-fna) is enabled.
+        backward_use_pt_reduction (bool): Whether to use PyTorch eager for computing the `dO * O`
+            product required by the backward pass, over the CUTLASS kernel. This only applies to
+            the `"cutlass-fmha"` backend.
+        run_persistent_kernel (bool): Whether to use persistent tile scheduling in the forward pass
+            kernel. This only applies to the `"blackwell-fmha"` backend.
+        kernel_schedule (Optional[str]): Kernel type (Hopper architecture only). Choices are
+            `None`: pick the default, `"non"` (non-persistent), `"coop"` (warp-specialized
+            cooperative), or `"pp"` (warp-specialized ping-ponging). Refer to
+            [Hopper FMHA/FNA backend](backends.md#hopper-fna-fmha) for more information.
+        torch_compile (bool): Applies only to the `"flex-fmha"` backend. Whether or not to JIT
+            compile the attention kernel. Due to this being an experimental feature in PyTorch, we
+            do not recommend it, and it is guarded by context flags. Read more in
+            [Flex Attention + `torch.compile`](context.md#flex-attention-torchcompile).
+        return_lse (bool): Whether or not to return the `logsumexp` tensor. `logsumexp` can be used
+            in the backward pass, and for [attention merging][natten.merge_attentions].
+    Returns:
+        output (Tensor): 4-D output tensor, with the heads last layout
+            (`[batch, seqlen, heads, head_dim_v]`).
+        logsumexp (Tensor): only returned when `return_lse=True`. 3-D logsumexp tensor, with the
+            heads last layout (`[batch, seqlen, heads]`).
+    """
+    fmha_tensor_checks(query, key, value)
+    (
+        cumulative_seqlen_Q,
+        cumulative_seqlen_KV,
+        max_seqlen_Q,
+        max_seqlen_KV,
+    ) = varlen_tensor_checks(
+        query=query,
+        key=key,
+        value=value,
+        seqlens_Q=seqlens_Q,
+        seqlens_KV=seqlens_KV,
+        cumulative_seqlen_Q=cumulative_seqlen_Q,
+        cumulative_seqlen_KV=cumulative_seqlen_KV,
+        max_seqlen_Q=max_seqlen_Q,
+        max_seqlen_KV=max_seqlen_KV,
+    )
+    is_varlen = cumulative_seqlen_Q is not None
+    scale = scale or query.shape[-1] ** -0.5
+    kernel_schedule = check_kernel_schedule(kernel_schedule)
+    backend = backend or choose_fmha_backend(
+        query,
+        key,
+        value,
+        is_causal=is_causal,
+        is_varlen=is_varlen,
+        torch_compile=torch_compile,
+    )
+    if backend == "blackwell-fmha":
+        return cutlass_blackwell_fmha(
+            query=query,
+            key=key,
+            value=value,
+            is_causal=is_causal,
+            scale=scale,
+            cumulative_seqlen_Q=cumulative_seqlen_Q,
+            cumulative_seqlen_KV=cumulative_seqlen_KV,
+            max_seqlen_Q=max_seqlen_Q,
+            max_seqlen_KV=max_seqlen_KV,
+            q_tile_size=q_tile_size,
+            kv_tile_size=kv_tile_size,
+            backward_q_tile_size=backward_q_tile_size,
+            backward_kv_tile_size=backward_kv_tile_size,
+            run_persistent_kernel=run_persistent_kernel,
+            return_lse=return_lse,
+        )
+    if backend == "hopper-fmha":
+        return cutlass_hopper_fmha(
+            query=query,
+            key=key,
+            value=value,
+            is_causal=is_causal,
+            scale=scale,
+            q_tile_size=q_tile_size,
+            kv_tile_size=kv_tile_size,
+            backward_q_tile_size=backward_q_tile_size,
+            backward_kv_tile_size=backward_kv_tile_size,
+            kernel_schedule=kernel_schedule,
+            return_lse=return_lse,
+            cumulative_seqlen_Q=cumulative_seqlen_Q,
+            cumulative_seqlen_KV=cumulative_seqlen_KV,
+            max_seqlen_Q=max_seqlen_Q,
+            max_seqlen_KV=max_seqlen_KV,
+        )
+    elif backend == "cutlass-fmha":
+        return cutlass_fmha(
+            query=query,
+            key=key,
+            value=value,
+            is_causal=is_causal,
+            scale=scale,
+            q_tile_size=q_tile_size,
+            kv_tile_size=kv_tile_size,
+            backward_q_tile_size=backward_q_tile_size,
+            backward_kv_tile_size=backward_kv_tile_size,
+            backward_kv_splits=backward_kv_splits,
+            backward_use_pt_reduction=backward_use_pt_reduction,
+            return_lse=return_lse,
+            cumulative_seqlen_Q=cumulative_seqlen_Q,
+            cumulative_seqlen_KV=cumulative_seqlen_KV,
+            max_seqlen_Q=max_seqlen_Q,
+            max_seqlen_KV=max_seqlen_KV,
+        )
+    elif backend == "flex-fmha":
+        return flex_fmha(
+            query=query,
+            key=key,
+            value=value,
+            is_causal=is_causal,
+            scale=scale,
+            q_tile_size=q_tile_size,
+            kv_tile_size=kv_tile_size,
+            torch_compile=torch_compile,
+            return_lse=return_lse,
+            cumulative_seqlen_Q=cumulative_seqlen_Q,
+            cumulative_seqlen_KV=cumulative_seqlen_KV,
+            max_seqlen_Q=max_seqlen_Q,
+            max_seqlen_KV=max_seqlen_KV,
+        )
+    raise NotImplementedError(f"Unrecognized NATTEN FMHA backend {backend}.")
+# Neighborhood Attention
+def neighborhood_attention_generic(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    kernel_size: DimensionTypeOrDed,
+    stride: DimensionTypeOrDed = 1,
+    dilation: DimensionTypeOrDed = 1,
+    is_causal: Optional[CausalArgTypeOrDed] = False,
+    scale: Optional[float] = None,
+    additional_keys: Optional[Tensor] = None,
+    additional_values: Optional[Tensor] = None,
+    attention_kwargs: Optional[Dict] = None,
+    # Perf-related args
+    backend: Optional[str] = None,
+    q_tile_shape: Optional[DimensionType] = None,
+    kv_tile_shape: Optional[DimensionType] = None,
+    backward_q_tile_shape: Optional[DimensionType] = None,
+    backward_kv_tile_shape: Optional[DimensionType] = None,
+    backward_kv_splits: Optional[DimensionType] = None,
+    backward_use_pt_reduction: bool = False,
+    run_persistent_kernel: bool = True,
+    kernel_schedule: Optional[Union[str, KernelSchedule]] = None,
+    torch_compile: bool = False,
+    return_lse: bool = False,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    na_tensor_checks(query, key, value)
+    additional_kv_tensor_checks(query, key, value, additional_keys, additional_values)
+    kernel_schedule = check_kernel_schedule(kernel_schedule)
+    na_dim = query.dim() - 3  # batch, heads, head_dim
+    assert na_dim in [1, 2, 3]
+    kernel_size, stride, dilation, is_causal = check_all_args(
+        na_dim, kernel_size, stride, dilation, is_causal
+    )
+    check_args_against_input(
+        query,
+        kernel_size=kernel_size,
+        stride=stride,
+        dilation=dilation,
+        is_causal=is_causal,
+    )
+    has_additional_attention = (
+        additional_keys is not None and additional_values is not None
+    )
+    if is_self_attention(
+        query,
+        kernel_size=kernel_size,
+        is_causal=is_causal,
+        has_additional_attention=has_additional_attention,
+    ):
+        logger.debug(
+            f"{query.shape=} with {kernel_size=}, {has_additional_attention=} and {is_causal=} is "
+            "self attention. Calling attention instead of neighborhood attention directly."
+        )
+        query_shape = query.shape
+        query = query.flatten(1, na_dim)
+        key = key.flatten(1, na_dim)
+        value = value.flatten(1, na_dim)
+        if has_additional_attention:
+            assert additional_keys is not None
+            assert additional_values is not None
+            key = torch.cat([key, additional_keys], dim=1)
+            value = torch.cat([value, additional_values], dim=1)
+        attn_kwargs = attention_kwargs or {}
+        out, lse = attention(
+            query,
+            key,
+            value,
+            is_causal=is_causal[0],  # NOTE: special case
+            scale=scale,
+            return_lse=True,
+            **attn_kwargs,
+        )
+        lse_shape = [s for s in query_shape[:-1]]
+        output_shape = lse_shape + [value.shape[-1]]
+        out = out.reshape(*output_shape)
+        lse = lse.reshape(*lse_shape)
+        if return_lse:
+            return out, lse
+        return out
+    scale = scale or query.shape[-1] ** -0.5
+    backend = backend or choose_backend(query, key, value, torch_compile=torch_compile)
+    if backend == "blackwell-fna":
+        output, lse = cutlass_blackwell_fna_generic(
+            query=query,
+            key=key,
+            value=value,
+            kernel_size=kernel_size,
+            stride=stride,
+            dilation=dilation,
+            is_causal=is_causal,
+            scale=scale,
+            q_tile_shape=q_tile_shape,
+            kv_tile_shape=kv_tile_shape,
+            backward_q_tile_shape=backward_q_tile_shape,
+            backward_kv_tile_shape=backward_kv_tile_shape,
+            run_persistent_kernel=run_persistent_kernel,
+            return_lse=True,
+        )
+    elif backend == "hopper-fna":
+        output, lse = cutlass_hopper_fna_generic(
+            query=query,
+            key=key,
+            value=value,
+            kernel_size=kernel_size,
+            stride=stride,
+            dilation=dilation,
+            is_causal=is_causal,
+            scale=scale,
+            q_tile_shape=q_tile_shape,
+            kv_tile_shape=kv_tile_shape,
+            backward_q_tile_shape=backward_q_tile_shape,
+            backward_kv_tile_shape=backward_kv_tile_shape,
+            kernel_schedule=kernel_schedule,
+            return_lse=True,
+        )
+    elif backend == "cutlass-fna":
+        output, lse = cutlass_fna_generic(
+            query=query,
+            key=key,
+            value=value,
+            kernel_size=kernel_size,
+            stride=stride,
+            dilation=dilation,
+            is_causal=is_causal,
+            scale=scale,
+            q_tile_shape=q_tile_shape,
+            kv_tile_shape=kv_tile_shape,
+            backward_q_tile_shape=backward_q_tile_shape,
+            backward_kv_tile_shape=backward_kv_tile_shape,
+            backward_kv_splits=backward_kv_splits,
+            backward_use_pt_reduction=backward_use_pt_reduction,
+            return_lse=True,
+        )
+    elif backend == "flex-fna":
+        output, lse = flex_fna_generic(
+            query=query,
+            key=key,
+            value=value,
+            kernel_size=kernel_size,
+            stride=stride,
+            dilation=dilation,
+            is_causal=is_causal,
+            scale=scale,
+            q_tile_shape=q_tile_shape,
+            kv_tile_shape=kv_tile_shape,
+            torch_compile=torch_compile,
+            return_lse=True,
+        )
+    else:
+        raise NotImplementedError(f"Unrecognized NATTEN backend {backend}.")
+    if has_additional_attention:
+        assert additional_keys is not None
+        assert additional_values is not None
+        attention_kwargs = attention_kwargs or {}
+        if "torch_compile" not in attention_kwargs:
+            attention_kwargs["torch_compile"] = torch_compile
+        additional_output, additional_lse = attention(
+            query.flatten(1, na_dim),
+            additional_keys,
+            additional_values,
+            is_causal=False,
+            scale=scale,
+            return_lse=True,
+            **attention_kwargs,
+        )
+        # NOTE: Flex unfused should not use the autograd fix
+        is_flex = backend == "flex-fna" or (
+            "backend" in attention_kwargs and attention_kwargs["backend"] == "flex-fmha"
+        )
+        use_autograd_fix = not is_flex or torch_compile
+        merged_output, merged_lse = merge_attentions(
+            [output.flatten(1, na_dim), additional_output],
+            [lse.flatten(1, na_dim), additional_lse],
+            use_autograd_fix=use_autograd_fix,
+        )
+        merged_output = merged_output.reshape(output.shape)
+        merged_lse = merged_lse.reshape(output.shape[:-1])
+        if return_lse:
+            return merged_output, merged_lse
+        return merged_output
+    if return_lse:
+        return output, lse
+    return output
+def na1d(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    kernel_size: Dimension1DTypeOrDed,
+    stride: Dimension1DTypeOrDed = 1,
+    dilation: Dimension1DTypeOrDed = 1,
+    is_causal: Optional[CausalArg1DTypeOrDed] = False,
+    scale: Optional[float] = None,
+    additional_keys: Optional[Tensor] = None,
+    additional_values: Optional[Tensor] = None,
+    attention_kwargs: Optional[Dict] = None,
+    backend: Optional[str] = None,
+    q_tile_shape: Optional[Dimension1DType] = None,
+    kv_tile_shape: Optional[Dimension1DType] = None,
+    backward_q_tile_shape: Optional[Dimension1DType] = None,
+    backward_kv_tile_shape: Optional[Dimension1DType] = None,
+    backward_kv_splits: Optional[Dimension1DType] = None,
+    backward_use_pt_reduction: bool = False,
+    run_persistent_kernel: bool = True,
+    kernel_schedule: Optional[Union[str, KernelSchedule]] = None,
+    torch_compile: bool = False,
+    return_lse: bool = False,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    """Computes 1-D neighborhood attention.
+    GQA/MQA support (`heads != heads_kv`) is available. For now, `blackwell-fna` and
+    `flex-fna` support GQA/MQA natively, and `cutlass-fna` and `hopper-fna` support it with
+    explicit repeats (increases memory usage and runtime).
+    Parameters:
+        query (Tensor): 4-D query tensor, with the heads last layout
+            (`[batch, seqlen, heads, head_dim]`)
+        key (Tensor): 4-D key tensor, with the heads last layout
+            (`[batch, seqlen, heads_kv, head_dim]`)
+        value (Tensor): 4-D value tensor, with the heads last layout
+            (`[batch, seqlen, heads_kv, head_dim_v]`)
+        kernel_size (Tuple[int] | int): Neighborhood window (kernel) size.
+            !!! note
+                `kernel_size` must be smaller than or equal to `seqlen`.
+        stride (Tuple[int] | int): Sliding window step size. Defaults to `1` (standard sliding
+            window).
+            !!! note
+                `stride` must be smaller than or equal to `kernel_size`.
+                When `stride == kernel_size`, there will be no overlap between sliding windows,
+                which is equivalent to blocked attention (a.k.a.
+                [window self attention](https://arxiv.org/abs/2103.14030)).
+        dilation (Tuple[int] | int): Dilation step size. Defaults to `1` (standard sliding window).
+            !!! note
+                The product of `dilation` and `kernel_size` must be smaller than or equal to
+                `seqlen`.
+        is_causal (Tuple[bool] | bool): Toggle causal masking. Defaults to `False`
+            (bi-directional).
+        scale (float): Attention scale. Defaults to `head_dim ** -0.5`.
+        additional_keys: `None` or 4-D key tensor, with the heads last layout
+            (`[batch, seqlen_kv, heads_kv, head_dim]`), corresponding to key tokens from some
+            additional context. Used when performing neighborhood cross-attention, where `query`
+            tokens attend to their neighborhood, as well as some fixed additional set of tokens.
+        additional_values: `None` or 4-D value tensor, with the heads last layout
+            (`[batch, seqlen_kv, heads_kv, head_dim_v]`), corresponding to value tokens from some
+            additional context. Used when performing neighborhood cross-attention, where `query`
+            tokens attend to their neighborhood, as well as some fixed additional set of tokens.
+            !!! note
+                `additional_keys` and `additional_values` must both either be `Tensor`s, or both
+                `None`s, and must match in shape.
+    Other Parameters:
+        backend (str): Backend implementation to run with. Choices are: `None` (pick the best
+            available one), `"cutlass-fna"`, `"hopper-fna"`, `"blackwell-fna"`, `"flex-fna"`.
+            Refer to [backends](backends.md) for more information.
+        q_tile_shape (Tuple[int]): 1-D Tile shape for the query token layout in the forward pass
+            kernel. You can use [profiler](profiler.md) to find valid choices for your use case,
+            and search for the best combination.
+        kv_tile_shape (Tuple[int]): 1-D Tile shape for the key-value token layout in the forward
+            pass kernel. You can use [profiler](profiler.md) to find valid choices for your use
+            case, and search for the best combination.
+        backward_q_tile_shape (Tuple[int]): 1-D Tile shape for the query token layout in the
+            backward pass kernel. This is ignored by `"flex-fna"`.
+            You can use [profiler](profiler.md) to find valid choices for your use case, and
+            search for the best combination.
+        backward_kv_tile_shape (Tuple[int]): 1-D Tile shape for the key/value token layout in the
+            backward pass kernel. This is ignored by `"flex-fna"`.
+            You can use [profiler](profiler.md) to find valid choices for your use case, and
+            search for the best combination.
+        backward_kv_splits (Tuple[int]): Number of key/value tiles allowed to work in parallel in
+            the backward pass kernel. Like tile shapes, this is a tuple and not an integer for
+            neighborhood attention operations, and the size of the tuple corresponds to the number
+            of dimensions / rank of the layout of tokens. This is only respected by the
+            `"cutlass-fna"` backend, and only when
+            [KV parallelism](context.md#kv-parallelism-in-fna) is enabled.
+        backward_use_pt_reduction (bool): Whether to use PyTorch eager for computing the `dO * O`
+            product required by the backward pass, over the CUTLASS kernel. This only applies to
+            the `"cutlass-fna"` backend.
+        run_persistent_kernel (bool): Whether to use persistent tile scheduling in the forward pass
+            kernel. This only applies to the `"blackwell-fna"` backend.
+        kernel_schedule (Optional[str]): Kernel type (Hopper architecture only). Choices are
+            `None`: pick the default, `"non"` (non-persistent), `"coop"` (warp-specialized
+            cooperative), or `"pp"` (warp-specialized ping-ponging). Refer to
+            [Hopper FMHA/FNA backend](backends.md#hopper-fna-fmha) for more information.
+        torch_compile (bool): Applies only to the `"flex-fna"` backend. Whether or not to JIT
+            compile the attention kernel. Due to this being an experimental feature in PyTorch, we
+            do not recommend it, and it is guarded by context flags. Read more in
+            [Flex Attention + `torch.compile`](context.md#flex-attention-torchcompile).
+        attention_kwargs: arguments to the [attention][natten.attention] operator, if used to
+            implement neighborhood cross-attention, or self attention as a fast path for
+            neighborhood attention.
+            If `additional_{keys,values}` are specified, NATTEN usually performs a separate
+            cross-attention using our [attention][natten.attention] operator, and
+            [merges][natten.merge_attentions] the results.
+            If for a given use case, the neighborhood attention problem is equivalent to self
+            attention (not causal, `kernel_size == seqlen`), NATTEN will also attempt to directly
+            use [attention][natten.attention].
+            You can override arguments to [attention][natten.attention] by passing a
+            dictionary here.
+            !!! example
+                ```python
+                out = na1d(
+                    q, k, v, kernel_size=kernel_size,
+                    ...,
+                    attention_kwargs={
+                        "backend": "blackwell-fmha",
+                        "run_persistent_kernel": True,
+                    }
+                )
+                ```
+        return_lse (bool): Whether or not to return the `logsumexp` tensor. `logsumexp` can be used
+            in the backward pass, and for [attention merging][natten.merge_attentions].
+    Returns:
+        output (Tensor): 4-D output tensor, with the heads last layout
+            (`[batch, seqlen, heads, head_dim_v]`).
+        logsumexp (Tensor): only returned when `return_lse=True`. 3-D logsumexp tensor, with the
+            heads last layout (`[batch, seqlen, heads]`).
+    """
+    return neighborhood_attention_generic(
+        query=query,
+        key=key,
+        value=value,
+        kernel_size=kernel_size,
+        stride=stride,
+        dilation=dilation,
+        is_causal=is_causal,
+        scale=scale,
+        additional_keys=additional_keys,
+        additional_values=additional_values,
+        attention_kwargs=attention_kwargs,
+        backend=backend,
+        q_tile_shape=q_tile_shape,
+        kv_tile_shape=kv_tile_shape,
+        backward_q_tile_shape=backward_q_tile_shape,
+        backward_kv_tile_shape=backward_kv_tile_shape,
+        backward_kv_splits=backward_kv_splits,
+        backward_use_pt_reduction=backward_use_pt_reduction,
+        run_persistent_kernel=run_persistent_kernel,
+        kernel_schedule=kernel_schedule,
+        torch_compile=torch_compile,
+        return_lse=return_lse,
+    )
+def na2d(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    kernel_size: Dimension2DTypeOrDed,
+    stride: Dimension2DTypeOrDed = 1,
+    dilation: Dimension2DTypeOrDed = 1,
+    is_causal: Optional[CausalArg2DTypeOrDed] = False,
+    scale: Optional[float] = None,
+    additional_keys: Optional[Tensor] = None,
+    additional_values: Optional[Tensor] = None,
+    attention_kwargs: Optional[Dict] = None,
+    backend: Optional[str] = None,
+    q_tile_shape: Optional[Dimension2DType] = None,
+    kv_tile_shape: Optional[Dimension2DType] = None,
+    backward_q_tile_shape: Optional[Dimension2DType] = None,
+    backward_kv_tile_shape: Optional[Dimension2DType] = None,
+    backward_kv_splits: Optional[Dimension2DType] = None,
+    backward_use_pt_reduction: bool = False,
+    run_persistent_kernel: bool = True,
+    kernel_schedule: Optional[Union[str, KernelSchedule]] = None,
+    torch_compile: bool = False,
+    return_lse: bool = False,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    """Computes 2-D neighborhood attention.
+    GQA/MQA support (`heads != heads_kv`) is available. For now, `blackwell-fna` and
+    `flex-fna` support GQA/MQA natively, and `cutlass-fna` and `hopper-fna` support it with
+    explicit repeats (increases memory usage and runtime).
+    Parameters:
+        query (Tensor): 2-D query tensor, with the heads last layout:
+            `[batch, X, Y, heads, head_dim]`, where token layout shape (feature map shape) is
+            `(X, Y)`.
+        key (Tensor): 2-D key tensor, with the heads last layout:
+            `[batch, X, Y, heads_kv, head_dim]`, where token layout shape (feature map shape) is
+            `(X, Y)`.
+        value (Tensor): 2-D value tensor, with the heads last layout:
+            `[batch, X, Y, heads_kv, head_dim_v]`, where token layout shape (feature map shape) is
+            `(X, Y)`.
+        kernel_size (Tuple[int, int] | int): Neighborhood window (kernel) size/shape. If an
+            integer, it will be repeated for all 2 dimensions. For example `kernel_size=3` is
+            reinterpreted as `kernel_size=(3, 3)`.
+            !!! note
+                `kernel_size` must be smaller than or equal to token layout shape (`(X, Y)`) along
+                every dimension.
+        stride (Tuple[int, int] | int): Sliding window step size/shape. Defaults to `1` (standard
+            sliding window). If an integer, it will be repeated for all 2 dimensions. For example
+            `stride=2` is reinterpreted as `stride=(2, 2)`.
+            !!! note
+                `stride` must be smaller than or equal to `kernel_size` along every dimension.
+                When `stride == kernel_size`, there will be no overlap between sliding windows,
+                which is equivalent to blocked attention (a.k.a.
+                [window self attention](https://arxiv.org/abs/2103.14030)).
+        dilation (Tuple[int, int] | int): Dilation step size/shape. Defaults to `1` (standard
+            sliding window). If an integer, it will be repeated for all 2 dimensions. For example
+            `dilation=4` is reinterpreted as `dilation=(4, 4)`.
+            !!! note
+                The product of `dilation` and `kernel_size` must be smaller than or equal to
+                token layout shape (`(X, Y)`) along every dimension.
+        is_causal (Tuple[bool, bool] | bool): Toggle causal masking. Defaults to `False`
+            (bi-directional). If a boolean, it will be repeated for all 2 dimensions. For example
+            `is_causal=True` is reinterpreted as `is_causal=(True, True)`.
+        scale (float): Attention scale. Defaults to `head_dim ** -0.5`.
+        additional_keys: `None` or 4-D key tensor, with the heads last layout
+            (`[batch, seqlen_kv, heads_kv, head_dim]`), corresponding to key tokens from some
+            additional context. Used when performing neighborhood cross-attention, where `query`
+            tokens attend to their neighborhood, as well as some fixed additional set of tokens.
+        additional_values: `None` or 4-D value tensor, with the heads last layout
+            (`[batch, seqlen_kv, heads_kv, head_dim_v]`), corresponding to value tokens from some
+            additional context. Used when performing neighborhood cross-attention, where `query`
+            tokens attend to their neighborhood, as well as some fixed additional set of tokens.
+            !!! note
+                `additional_keys` and `additional_values` must both either be `Tensor`s, or both
+                `None`s, and must match in shape.
+    Other Parameters:
+        backend (str): Backend implementation to run with. Choices are: `None` (pick the best
+            available one), `"cutlass-fna"`, `"hopper-fna"`, `"blackwell-fna"`, `"flex-fna"`.
+            Refer to [backends](backends.md) for more information.
+        q_tile_shape (Tuple[int, int]): 2-D Tile shape for the query token layout in the forward
+            pass kernel. You can use [profiler](profiler.md) to find valid choices for your use
+            case, and search for the best combination.
+        kv_tile_shape (Tuple[int, int]): 2-D Tile shape for the key-value token layout in the
+            forward pass kernel. You can use [profiler](profiler.md) to find valid choices for your
+            use case, and search for the best combination.
+        backward_q_tile_shape (Tuple[int, int]): 2-D Tile shape for the query token layout in the
+            backward pass kernel. This is ignored by `"flex-fna"`.
+            You can use [profiler](profiler.md) to find valid choices for your use case, and
+            search for the best combination.
+        backward_kv_tile_shape (Tuple[int, int]): 2-D Tile shape for the key/value token layout in
+            the backward pass kernel. This is ignored by `"flex-fna"`.
+            You can use [profiler](profiler.md) to find valid choices for your use case, and
+            search for the best combination.
+        backward_kv_splits (Tuple[int, int]): Number of key/value tiles allowed to work in parallel
+            in the backward pass kernel. Like tile shapes, this is a tuple and not an integer for
+            neighborhood attention operations, and the size of the tuple corresponds to the number
+            of dimensions / rank of the layout of tokens. This is only respected by the
+            `"cutlass-fna"` backend, and only when
+            [KV parallelism](context.md#kv-parallelism-in-fna) is enabled.
+        backward_use_pt_reduction (bool): Whether to use PyTorch eager for computing the `dO * O`
+            product required by the backward pass, over the CUTLASS kernel. This only applies to
+            the `"cutlass-fna"` backend.
+        run_persistent_kernel (bool): Whether to use persistent tile scheduling in the forward pass
+            kernel. This only applies to the `"blackwell-fna"` backend.
+        kernel_schedule (Optional[str]): Kernel type (Hopper architecture only). Choices are
+            `None`: pick the default, `"non"` (non-persistent), `"coop"` (warp-specialized
+            cooperative), or `"pp"` (warp-specialized ping-ponging). Refer to
+            [Hopper FMHA/FNA backend](backends.md#hopper-fna-fmha) for more information.
+        torch_compile (bool): Applies only to the `"flex-fna"` backend. Whether or not to JIT
+            compile the attention kernel. Due to this being an experimental feature in PyTorch, we
+            do not recommend it, and it is guarded by context flags. Read more in
+            [Flex Attention + `torch.compile`](context.md#flex-attention-torchcompile).
+        attention_kwargs: arguments to the [attention][natten.attention] operator, if used to
+            implement neighborhood cross-attention, or self attention as a fast path for
+            neighborhood attention.
+            If `additional_{keys,values}` are specified, NATTEN usually performs a separate
+            cross-attention using our [attention][natten.attention] operator, and
+            [merges][natten.merge_attentions] the results.
+            If for a given use case, the neighborhood attention problem is equivalent to self
+            attention (not causal along any dims, `kernel_size == (X, Y)`), NATTEN will also
+            attempt to directly use [attention][natten.attention].
+            You can override arguments to [attention][natten.attention] by passing a
+            dictionary here.
+            !!! example
+                ```python
+                out = na2d(
+                    q, k, v, kernel_size=kernel_size,
+                    ...,
+                    attention_kwargs={
+                        "backend": "blackwell-fmha",
+                        "run_persistent_kernel": True,
+                    }
+                )
+                ```
+        return_lse (bool): Whether or not to return the `logsumexp` tensor. `logsumexp` can be used
+            in the backward pass, and for [attention merging][natten.merge_attentions].
+    Returns:
+        output (Tensor): 5-D output tensor, with the heads last layout
+            (`[batch, X, Y, heads, head_dim_v]`).
+        logsumexp (Tensor): only returned when `return_lse=True`. 4-D logsumexp tensor, with the
+            heads last layout (`[batch, X, Y, heads]`).
+    """
+    return neighborhood_attention_generic(
+        query=query,
+        key=key,
+        value=value,
+        kernel_size=kernel_size,
+        stride=stride,
+        dilation=dilation,
+        is_causal=is_causal,
+        scale=scale,
+        additional_keys=additional_keys,
+        additional_values=additional_values,
+        attention_kwargs=attention_kwargs,
+        backend=backend,
+        q_tile_shape=q_tile_shape,
+        kv_tile_shape=kv_tile_shape,
+        backward_q_tile_shape=backward_q_tile_shape,
+        backward_kv_tile_shape=backward_kv_tile_shape,
+        backward_kv_splits=backward_kv_splits,
+        backward_use_pt_reduction=backward_use_pt_reduction,
+        run_persistent_kernel=run_persistent_kernel,
+        kernel_schedule=kernel_schedule,
+        torch_compile=torch_compile,
+        return_lse=return_lse,
+    )
+def na3d(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    kernel_size: Dimension3DTypeOrDed,
+    stride: Dimension3DTypeOrDed = 1,
+    dilation: Dimension3DTypeOrDed = 1,
+    is_causal: Optional[CausalArg3DTypeOrDed] = False,
+    scale: Optional[float] = None,
+    additional_keys: Optional[Tensor] = None,
+    additional_values: Optional[Tensor] = None,
+    attention_kwargs: Optional[Dict] = None,
+    backend: Optional[str] = None,
+    q_tile_shape: Optional[Dimension3DType] = None,
+    kv_tile_shape: Optional[Dimension3DType] = None,
+    backward_q_tile_shape: Optional[Dimension3DType] = None,
+    backward_kv_tile_shape: Optional[Dimension3DType] = None,
+    backward_kv_splits: Optional[Dimension3DType] = None,
+    backward_use_pt_reduction: bool = False,
+    run_persistent_kernel: bool = True,
+    kernel_schedule: Optional[Union[str, KernelSchedule]] = None,
+    torch_compile: bool = False,
+    return_lse: bool = False,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    """Computes 3-D neighborhood attention.
+    GQA/MQA support (`heads != heads_kv`) is available. For now, `blackwell-fna` and
+    `flex-fna` support GQA/MQA natively, and `cutlass-fna` and `hopper-fna` support it with
+    explicit repeats (increases memory usage and runtime).
+    Parameters:
+        query (Tensor): 3-D query tensor, with the heads last layout:
+            `[batch, X, Y, Z, heads, head_dim]`, where token layout shape (feature map shape) is
+            `(X, Y, Z)`.
+        key (Tensor): 3-D key tensor, with the heads last layout:
+            `[batch, X, Y, Z, heads_kv, head_dim]`, where token layout shape (feature map shape) is
+            `(X, Y, Z)`.
+        value (Tensor): 3-D value tensor, with the heads last layout:
+            `[batch, X, Y, Z, heads_kv, head_dim_V]`, where token layout shape (feature map shape) is
+            `(X, Y, Z)`.
+        kernel_size (Tuple[int, int, int] | int): Neighborhood window (kernel) size/shape. If an
+            integer, it will be repeated for all 3 dimensions. For example `kernel_size=3` is
+            reinterpreted as `kernel_size=(3, 3, 3)`.
+            !!! note
+                `kernel_size` must be smaller than or equal to token layout shape (`(X, Y, Z)`)
+                along every dimension.
+        stride (Tuple[int, int, int] | int): Sliding window step size/shape. Defaults to `1`
+            (standard sliding window). If an integer, it will be repeated for all 3 dimensions.
+            For example `stride=2` is reinterpreted as `stride=(2, 2, 2)`.
+            !!! note
+                `stride` must be smaller than or equal to `kernel_size` along every dimension.
+                When `stride == kernel_size`, there will be no overlap between sliding windows,
+                which is equivalent to blocked attention (a.k.a.
+                [window self attention](https://arxiv.org/abs/2103.14030)).
+        dilation (Tuple[int, int, int] | int): Dilation step size/shape. Defaults to `1` (standard
+            sliding window). If an integer, it will be repeated for all 3 dimensions. For example
+            `dilation=4` is reinterpreted as `dilation=(4, 4, 4)`.
+            !!! note
+                The product of `dilation` and `kernel_size` must be smaller than or equal to
+                token layout shape (`(X, Y, Z)`) along every dimension.
+        is_causal (Tuple[bool, bool, bool] | bool): Toggle causal masking. Defaults to `False`
+            (bi-directional). If a boolean, it will be repeated for all 3 dimensions. For example
+            `is_causal=True` is reinterpreted as `is_causal=(True, True, True)`.
+        scale (float): Attention scale. Defaults to `head_dim ** -0.5`.
+        additional_keys: `None` or 4-D key tensor, with the heads last layout
+            (`[batch, seqlen_kv, heads_kv, head_dim_v]`), corresponding to key tokens from some
+            additional context. Used when performing neighborhood cross-attention, where `query`
+            tokens attend to their neighborhood, as well as some fixed additional set of tokens.
+        additional_values: `None` or 4-D value tensor, with the heads last layout
+            (`[batch, seqlen_kv, heads_kv, head_dim_v]`), corresponding to value tokens from some
+            additional context. Used when performing neighborhood cross-attention, where `query`
+            tokens attend to their neighborhood, as well as some fixed additional set of tokens.
+            !!! note
+                `additional_keys` and `additional_values` must both either be `Tensor`s, or both
+                `None`s, and must match in shape.
+    Other Parameters:
+        backend (str): Backend implementation to run with. Choices are: `None` (pick the best
+            available one), `"cutlass-fna"`, `"hopper-fna"`, `"blackwell-fna"`, `"flex-fna"`.
+            Refer to [backends](backends.md) for more information.
+        q_tile_shape (Tuple[int, int, int]): 3-D Tile shape for the query token layout in the
+            forward pass kernel. You can use [profiler](profiler.md) to find valid choices for your
+            use case, and search for the best combination.
+        kv_tile_shape (Tuple[int, int, int]): 3-D Tile shape for the key-value token layout in the
+            forward pass kernel. You can use [profiler](profiler.md) to find valid choices for your
+            use case, and search for the best combination.
+        backward_q_tile_shape (Tuple[int, int, int]): 3-D Tile shape for the query token layout in
+            the backward pass kernel. This is ignored by `"flex-fna"`.
+            You can use [profiler](profiler.md) to find valid choices for your use case, and
+            search for the best combination.
+        backward_kv_tile_shape (Tuple[int, int, int]): 3-D Tile shape for the key/value token
+            layout in the backward pass kernel. This is ignored by `"flex-fna"`.
+            You can use [profiler](profiler.md) to find valid choices for your use case,
+            and search for the best combination.
+        backward_kv_splits (Tuple[int, int, int]): Number of key/value tiles allowed to work in
+            parallel in the backward pass kernel. Like tile shapes, this is a tuple and not an
+            integer for neighborhood attention operations, and the size of the tuple corresponds to
+            the number of dimensions / rank of the layout of tokens. This is only respected by the
+            `"cutlass-fna"` backend, and only when
+            [KV parallelism](context.md#kv-parallelism-in-fna) is enabled.
+        backward_use_pt_reduction (bool): Whether to use PyTorch eager for computing the `dO * O`
+            product required by the backward pass, over the CUTLASS kernel. This only applies to
+            the `"cutlass-fna"` backend.
+        run_persistent_kernel (bool): Whether to use persistent tile scheduling in the forward pass
+            kernel. This only applies to the `"blackwell-fna"` backend.
+        kernel_schedule (Optional[str]): Kernel type (Hopper architecture only). Choices are
+            `None`: pick the default, `"non"` (non-persistent), `"coop"` (warp-specialized
+            cooperative), or `"pp"` (warp-specialized ping-ponging). Refer to
+            [Hopper FMHA/FNA backend](backends.md#hopper-fna-fmha) for more information.
+        torch_compile (bool): Applies only to the `"flex-fna"` backend. Whether or not to JIT
+            compile the attention kernel. Due to this being an experimental feature in PyTorch, we
+            do not recommend it, and it is guarded by context flags. Read more in
+            [Flex Attention + `torch.compile`](context.md#flex-attention-torchcompile).
+        attention_kwargs: arguments to the [attention][natten.attention] operator, if used to
+            implement neighborhood cross-attention, or self attention as a fast path for
+            neighborhood attention.
+            If `additional_{keys,values}` are specified, NATTEN usually performs a separate
+            cross-attention using our [attention][natten.attention] operator, and
+            [merges][natten.merge_attentions] the results.
+            If for a given use case, the neighborhood attention problem is equivalent to self
+            attention (not causal along any dims, `kernel_size == (X, Y, Z)`), NATTEN will also
+            attempt to directly use [attention][natten.attention].
+            You can override arguments to [attention][natten.attention] by passing a
+            dictionary here.
+            !!! example
+                ```python
+                out = na3d(
+                    q, k, v, kernel_size=kernel_size,
+                    ...,
+                    attention_kwargs={
+                        "backend": "blackwell-fmha",
+                        "run_persistent_kernel": True,
+                    }
+                )
+                ```
+        return_lse (bool): Whether or not to return the `logsumexp` tensor. `logsumexp` can be used
+            in the backward pass, and for [attention merging][natten.merge_attentions].
+    Returns:
+        output (Tensor): 6-D output tensor, with the heads last layout
+            (`[batch, X, Y, Z, heads, head_dim_v]`).
+        logsumexp (Tensor): only returned when `return_lse=True`. 5-D logsumexp tensor, with the
+            heads last layout (`[batch, X, Y, Z, heads]`).
+    """
+    return neighborhood_attention_generic(
+        query=query,
+        key=key,
+        value=value,
+        kernel_size=kernel_size,
+        stride=stride,
+        dilation=dilation,
+        is_causal=is_causal,
+        scale=scale,
+        additional_keys=additional_keys,
+        additional_values=additional_values,
+        attention_kwargs=attention_kwargs,
+        backend=backend,
+        q_tile_shape=q_tile_shape,
+        kv_tile_shape=kv_tile_shape,
+        backward_q_tile_shape=backward_q_tile_shape,
+        backward_kv_tile_shape=backward_kv_tile_shape,
+        backward_kv_splits=backward_kv_splits,
+        backward_use_pt_reduction=backward_use_pt_reduction,
+        run_persistent_kernel=run_persistent_kernel,
+        kernel_schedule=kernel_schedule,
+        torch_compile=torch_compile,
+        return_lse=return_lse,
+    )

build/torch210-cu126-aarch64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,84 @@

+{
+  "name": "natten",
+  "id": "_natten_cuda_5fe9dbb",
+  "version": 1,
+  "kernels-minver": "0.14.0",
+  "license": "MIT",
+  "upstream": "https://github.com/SHI-Labs/NATTEN",
+  "python-depends": [],
+  "kernel-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "8.0",
+      "9.0",
+      "9.0a"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "2xTdlXXAhz/SpvyCABtaEDtbZW55DRw4YDDMSF94Kp4=",
+      "_environment.py": "9P189XTk2YMyWuRNI5eCI3UvHRqVR8Zh2hgJsHP7MYI=",
+      "_libnatten/__init__.py": "yMW+2Kd5Nxs716jaNq15IleYQ86HR5u7SMEQDqVVunI=",
+      "_libnatten/torch_wrappers.py": "v7vkrytkghNz+cJI7l2PXGqWactsSM1ImwCsHfyQ3iA=",
+      "_natten_cuda_5fe9dbb.abi3.so": "AvnzzKD+lh9GDOhIAXzHo5GqdONMop5fZDHRE4xDNgA=",
+      "_ops.py": "MbEvw83543TKiDVD1FyrMuwUK7cOvCbAaagp4d2oEUk=",
+      "_private_for_testing.py": "Ir5abh5fQ1+9aqmmi7RqGhZybEvYV8JEnsrhzj558xw=",
+      "_types.py": "OzK1SfxBXbR01LDSzFOpvm5qUCLs3qOc021I31Eibm8=",
+      "attn_merge.py": "KDItp+YcD/18PEXqCOOxSBeLYYt50LhYZnyluE5iT7A=",
+      "backends/__init__.py": "r8Lco2ESXb2yqPccmIu6kyUD8niprXF7J0WA7urlAbE=",
+      "backends/blackwell_fmha.py": "hDQuJ27vACJ46dj4QKaa7CWB2Sy7650ZtHlJCgQQ3Bk=",
+      "backends/blackwell_fna.py": "YA6YqgDnTJ/4F5FCe7uf7H0qOMhF4k4UTkKVtYtfYDU=",
+      "backends/configs/__init__.py": "NG3ArX78+S0/E3i2R0LTjmr30ZY71qf9Xv3396Wfw34=",
+      "backends/configs/checks.py": "d7qCFUqvF9fl0nuRiTr6eDXyaAf0/ksyrkqI8Z2VG+E=",
+      "backends/configs/cutlass/__init__.py": "H+6PW4E9/I2UtU0CUB+8rEvSjpT6hqVHEjO5JQD0cao=",
+      "backends/configs/cutlass/backward_knobs.py": "s6CakBU+ApSjFfLGLyC42kpmXZBc2Si3Ud+4SKIZ7zM=",
+      "backends/configs/cutlass/fna_backward_128x128.py": "ZcUo6u0hDGUXasU52U98XV+KUF475QyP8I5O8T042yM=",
+      "backends/configs/cutlass/fna_backward_128x64.py": "3gXgi7PqQpP8A7ClS8GPA1eXwOiLCgPEeOL4oJxYWXY=",
+      "backends/configs/cutlass/fna_backward_64x64.py": "i5iRd+TrdtDS8S48cVj6/D7+951L6lPLNuqkLIA1k/c=",
+      "backends/configs/cutlass/fna_forward_32x128.py": "1acDc2mpRTKBC3opzjZEEr0R9lPmQ3wLPSOIlJlM5aM=",
+      "backends/configs/cutlass/fna_forward_64x128.py": "8NfkwqYqJ8FfdFeXZU1e6fUF8F6dRUZRsJ+c9aVDZAI=",
+      "backends/configs/cutlass/fna_forward_64x64.py": "vu1gBAvQLQuO+iJdz4GdidayIywG13b83Lorw/I43ac=",
+      "backends/configs/cutlass_blackwell/__init__.py": "RtGGNcP4bvzu1JozMLd16mJtDncxRIAa+W7Lg7I2KHs=",
+      "backends/configs/cutlass_hopper/__init__.py": "DjEcfThaXh2JGrA5EW72mmRdh6Kwk/oOmaOkZm9PkNU=",
+      "backends/configs/flex/__init__.py": "HFK8hxYrU85UyEZPO6c66DoLvfJcji6jmTNFUGzgGDs=",
+      "backends/flex.py": "6xSMb2kuezRaeOcBN4UWY8F6kJG7F6DoZdSv08vvmMg=",
+      "backends/fmha.py": "sA52KmY/Dj5Qm5E2eX9IWNVMke0Vn94FAU+tYe1zPLc=",
+      "backends/fna.py": "xUzyoiHkFwDyf+Cv0VOhDt+5X5F8X6ckVAo5cJZXGsc=",
+      "backends/hopper_fmha.py": "+xWkbfNibNwTtU8d9TmUUedITYnnUQwykYCgBcoyTkE=",
+      "backends/hopper_fna.py": "fKfBfO0gNJtdJeghLDa4tQpCdIeOvlkDLf5PaD0P2P8=",
+      "backends/reference.py": "SG0s45Zme6O8gSjPAxw3ib9T0pqvDAxCmbGn/DQ/0tc=",
+      "context.py": "IzU+TvVnBZ41jc93BjQ7+/Q3Eg+8fDNSK4YCbTc2OU8=",
+      "functional.py": "zQq6Et4t+ywceLDcrQmiXYh2xpkRqSk1vBdHE7QsCGE=",
+      "modules.py": "1C3KPU4tOHIuRT+rOd6XVqPt98Ho/oehFEKOpwv445Q=",
+      "symbols.json": "eyAwy4AHhZ5N1bQLcFp8sj715MIRx60EcLLk/lhcjKg=",
+      "token_permute/__init__.py": "NPxL8EMHxAxjHMcHTMFznH0oNQjGUr2divAbtGaWNhs=",
+      "token_permute/cutlass_impl.py": "nCXz+b8TJAbU/j0RiB0PMO94nqFZIvNxO9tKhLC/vDs=",
+      "token_permute/frontend.py": "mVlwIP2pjFEq/KHEiT9t9qE3Zl+pVpeeVhceclOIVOU=",
+      "token_permute/torch_impl.py": "7PKSTynMOiaEaGRxEoyWz6SVrC0Lcx227Cvqdc4kB9E=",
+      "utils/__init__.py": "+RppytrGJc5JC8/fPz8XWYdCwAT8JOxpmLDQTlpO6gI=",
+      "utils/checks.py": "FP+WlZSIT7sjyzKOjnytCuFbqN2PejuKJ6SyApkZPn0=",
+      "utils/device.py": "3UJXsbft+HlnP2+nAnIoE1Hlq5yq9LhLSNFx9I5QHaw=",
+      "utils/dtype.py": "p6d+m2q4qj480Kcqg/b7WeV8+n4oRr5vi+dMb9vsxE4=",
+      "utils/environment.py": "GyeBUYS5y09z0h50VEL09PWCutXI57kWGxnzX5GUvVM=",
+      "utils/log.py": "pvlKMKeLzqrNvjHgygV985sT0TjtnbYcjluTDJ4y2Mc=",
+      "utils/tensor.py": "PhXmULtQny1AVmZAMByMKV7xNm/gf6M6VQa3gkjQrTQ=",
+      "utils/testing.py": "nl8DQHWiR6AWmecX2mHpkZxZBomc9QRHXc1MqvF9AG8=",
+      "utils/tuples.py": "/LrawWeD5LDugY/SD92NNcr9M3BNfs/F+8lSYs+130Y=",
+      "utils/varlen.py": "ON5q6wFAyBOY0PscRYht7LdGrwr5p96TcojdrprVaHI=",
+      "version.py": "FVIe4O+OfUJ2NXT3Ebezok+DAPBcOkgEU6tlHu0PPg8="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.18.0-dev0",
+      "commit": "21bdf69d97a289b9c581a0c60f63deed131647d5",
+      "dirty": false
+    },
+    "kernel": {
+      "commit": "5fe9dbb3d45f6f2652958c1f70e5381891fa7aa0",
+      "dirty": false
+    }
+  }
+}

build/torch210-cu126-aarch64-linux/metadata.json.sigstore ADDED Viewed

	@@ -0,0 +1 @@

+ {"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHLDCCBrOgAwIBAgIUConfL87twl5y8Cnu/9yMUMstpzEwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDMxMDI1NzM3WhcNMjYxMDMxMDMwNzM3WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAE5OUqvvfMUNlTnXuzgspqC4+nmM9rjRCils9yzXXtR/Raf+lhIphGnxWm1FaoSMVuL9e6NYIxUfud8qpvORjyaaOCBdIwggXOMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBCg1ZmU5ZGJiM2Q0NWY2ZjI2NTI5NThjMWY3MGU1MzgxODkxZmE3YWEwMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDCg1ZmU5ZGJiM2Q0NWY2ZjI2NTI5NThjMWY3MGU1MzgxODkxZmE3YWEwMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDCg1ZmU5ZGJiM2Q0NWY2ZjI2NTI5NThjMWY3MGU1MzgxODkxZmE3YWEwMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoNWZlOWRiYjNkNDVmNmYyNjUyOTU4YzFmNzBlNTM4MTg5MWZhN2FhMDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy80MTA4MzU5MjA4MS9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGLBgorBgEEAdZ5AgQCBH0EewB5AHcA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhj+QqDgAABAMASDBGAiEAh4ZF6nzagTaObEPoqfafHCKvVKi8mJ/GqCTKt8o+RU0CIQD68C34HN171/SvQSu9LHuEh9BNb/j2WTE71CgYiHzy3TAKBggqhkjOPQQDAwNnADBkAjBU/vEGqs8gSmAx9y4nh4nOe8z/19Gx5NswnQrF1LmFqUP42Mot+KThfeQvX04zJ/MCMEYxcKQ9BN4KT5M2ppowh5fXPpjTlv11EVBTjXWOUa9+lmgag8eLTy7qERjPA2QwdQ=="},"tlogEntries":[{"logIndex":"3608411523","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1793415457","inclusionPromise":{"signedEntryTimestamp":"MEYCIQCUmIxOX/tqdAVOw5F2IiljP6VNOhlUttOb3chnkXlGuQIhAJ0qdYdqLF8FXYjptF8cDzWU8nkClAOiGuBHZRI+D8m3"},"inclusionProof":{"logIndex":"3486507261","rootHash":"rWaNS1EkW1wLyjBp5b43Nq/5vvPR6YcQQQGZDB6cQ8U=","treeSize":"3486507265","hashes":["QNaWrpma95Q2Ii9e7ecI6vH1kYwmw7kh0Tm8OM8EE34=","AFALTjX8vYkbGFpZGb5eNujiOLa479EoyhE8lsvBAqs=","Qh8cl2UbYM9EOL+CeWrW2WLvK6W1R+AKOk1abKiQor8=","8gpPATA8/zm5NyGa4N+3RZZFhLUdH0+ZpwotTIwcFYg=","9754Egf02ycYSNnC/A4rAfL956NfkfVXqOqnvdGmfF4=","fLvFj4rPJDFcCGPK+47k1APgaSzVu7kh+7/3nfcOhmc=","/6Y13fDlheN2bcnEsII/U9MAyKbIqLnFWAcVACDDQCw=","SlZV3uwFo6k72lrltHbOGXXvvNPan6CnD/J18uIb//c=","jyx4f5I0hMp434CJXP0Rut0PE282/+M4ouvnTUsV/UE=","J4FrfDFXm8W/ieq4gOlI681sJX3Ml/5l1aS1rO+w6mY=","3GcBUVQq6SethH9kno916KUFjBo/hYxuR4FwCIHSVM4=","h2QiGRcZnh9Vy0+1n4SpULgQJQzA7gslblU6vHJAgSw=","kvXR/O9NmwtrbMarzteT806WdkrAbUJZB1qow2SgxkQ=","3vFhaeVJmWIEPQcb814RjFRDT1tH7McB24OPJ32si80=","lE7/U/HjlRgHzdP8vPTBU4Zaz3UH2C7YdUvkQUuDfnQ=","Bz1DtqAVEHz4Bbo/X5KzdwHn3d91F2+g+KUOWnMj/2w=","C1iNsNunwIhtOOJvTDxee6pgbCkOjryTGo1g+5SLxEs=","2khvWt9d+ZdJVlsE9c5CUGSWz42nP08VUkB5e4aYzxI=","e3OaWwzL1tEWYXiUzGpH+1v9CiFPwweyfRzIh3Z+uO4=","b+938HIdfjdD/+LgIs3AVRE3Wc65/1NSSAL4OJ5FUXE=","kQz/4ND7LdhfAQn+idNsO+ZFC32sfkYRlhjWF5W+sPc=","zGfeoWKmrZlXMOvyi3R2GDbGpmI1skJW/wLTXCCgF0o=","1DaDKzrCNNBhmSdY5vFjz+NpAmxmYfYEgBCdPLrMP1A=","xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3486507265\nrWaNS1EkW1wLyjBp5b43Nq/5vvPR6YcQQQGZDB6cQ8U=\n\n— rekor.sigstore.dev wNI9ajBGAiEAzJ3kviUPOvuZ83hVohKhY6vxrvD5zMXOTyIEMA4+DV8CIQDN2A0Mn/e9QgijyyztwE+qiQIsfJCBGD3wtNM6beAX4Q==\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI5MDE4YmZkYmZkMzRmOTZmNmUwZjE5ZGRkOTM1NWRhZDViNGVlMTNjMzVkNTM3ODg3ZGE0YmZiOTVlZWMzMTEyIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUURTZmRNa2Z5UFJTbmFDRGhwOEg1QmxwdDVYNzFrMHpqZGdCVFpPQ2hYcUNnSWdQQjE3Ym1qZVEwWmdZTUNpTkV6TkpRM2JQeHBYS3JHMXdHQmtCQ0E2RmRBPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1SRU5EUW5KUFowRjNTVUpCWjBsVlEyOXVaa3c0TjNSM2JEVjVPRU51ZFM4NWVVMVZUWE4wY0hwRmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFVFhoTlJFa3hUbnBOTTFkb1kwNU5hbGw0VFVSTmVFMUVUWGRPZWswelYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVUxVDFWeGRuWm1UVlZPYkZSdVdIVjZaM053Y1VNMEsyNXRUVGx5YWxKRGFXeHpPWGtLZWxoWWRGSXZVbUZtSzJ4b1NYQm9SMjU0VjIweFJtRnZVMDFXZFV3NVpUWk9XVWw0VldaMVpEaHhjSFpQVW1wNVlXRlBRMEprU1hkbloxaFBUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmN4V20xVk5RcGFSMHBwVFRKUk1FNVhXVEphYWtreVRsUkpOVTVVYUdwTlYxa3pUVWRWTVUxNlozaFBSR3Q0V20xRk0xbFhSWGROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEWnpGYWJWVTFXa2RLYVUweVVUQk9WMWt5Q2xwcVNUSk9WRWsxVGxSb2FrMVhXVE5OUjFVeFRYcG5lRTlFYTNoYWJVVXpXVmRGZDAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEWnpGYWJWVTFDbHBIU21sTk1sRXdUbGRaTWxwcVNUSk9WRWsxVGxSb2FrMVhXVE5OUjFVeFRYcG5lRTlFYTNoYWJVVXpXVmRGZDAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlPVjFwc1QxZFNhVmxxVG10T1JGWnRUbTFaZVU1cVZYbFBWRlUwV1hwR2JVNTZRbXhPVkUwMFRWUm5OVTFYV21nS1RqSkdhRTFFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRnd1RWUkJORTE2VlRWTmFrRTBUVk01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTUNrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NEQkZaWGRDTlVGSVkwRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvYWl0UmNVUm5RVUZDUVUxQlUwUkNSMEZwUlVGb05GcEdObTU2WVdkVVlVOWlSVkJ2Y1daaFpraERTM1pXUzJrNGJVb3ZSd3B4UTFSTGREaHZLMUpWTUVOSlVVUTJPRU16TkVoT01UY3hMMU4yVVZOMU9VeElkVVZvT1VKT1lpOXFNbGRVUlRjeFEyZFphVWg2ZVROVVFVdENaMmR4Q21ocmFrOVFVVkZFUVhkT2JrRkVRbXRCYWtKVkwzWkZSM0Z6T0dkVGJVRjRPWGswYm1nMGJrOWxPSG92TVRsSGVEVk9jM2R1VVhKR01VeHRSbkZWVURRS01rMXZkQ3RMVkdobVpWRjJXREEwZWtvdlRVTk5SVmw0WTB0Uk9VSk9ORXRVTlUweWNIQnZkMmcxWmxoUWNHcFViSFl4TVVWV1FsUnFXRmRQVldFNUt3cHNiV2RoWnpobFRGUjVOM0ZGVW1wUVFUSlJkMlJSUFQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgeBtuKIWobPfafWLTeWzCcYQk473v8o1hFCxgUxHneBoCFQDWLvslBwhsQMCtkdTOoxfktdPljBgPMjAyNjEwMzEwMjU3MzdaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHaMIIB1gIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAzMTAyNTczN1owLwYJKoZIhvcNAQkEMSIEILUq8Bqu4E+u4OQxwFRxE7I5oDubJkPB4rsX2ZUoyjXlMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRmMGQCMHCbkV5CWTiOCvmEDQrDYS5FmkUxX4bFnx17paQEeX2hnzJK3SoUKlnbLo0KnqS7EQIwYW+fCUf4rgfYJBriGLKbx9HJUt8aaSNsuz+q2XEAu1pnl4Pgo9/L/2hf2vrGiHlU"}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"kBi/2/00+W9uDxnd2TVdrVtO4Tw11TeIfaS/uV7sMRI="},"signature":"MEUCIQDSfdMkfyPRSnaCDhp8H5Blpt5X71k0zjdgBTZOChXqCgIgPB17bmjeQ0ZgYMCiNEzNJQ3bPxpXKrG1wGBkBCA6FdA="}}

build/torch210-cu126-aarch64-linux/modules.py ADDED Viewed

	@@ -0,0 +1,449 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+from typing import Optional
+import torch  # noqa: F401
+from torch import nn, Tensor
+from .functional import neighborhood_attention_generic
+from ._types import (
+    CausalArg1DTypeOrDed,
+    CausalArg2DTypeOrDed,
+    CausalArg3DTypeOrDed,
+    CausalArgTypeOrDed,
+    Dimension1DTypeOrDed,
+    Dimension2DTypeOrDed,
+    Dimension3DTypeOrDed,
+    DimensionTypeOrDed,
+)
+from .utils.checks import check_all_args
+class NeighborhoodAttentionGeneric(nn.Module):
+    def __init__(
+        self,
+        na_dim: int,
+        embed_dim: int,
+        num_heads: int,
+        kernel_size: DimensionTypeOrDed,
+        stride: DimensionTypeOrDed = 1,
+        dilation: DimensionTypeOrDed = 1,
+        is_causal: CausalArgTypeOrDed = False,
+        qkv_bias: bool = True,
+        qk_scale: Optional[float] = None,
+        proj_drop: float = 0.0,
+    ):
+        super().__init__()
+        kernel_size, stride, dilation, is_causal = check_all_args(
+            na_dim, kernel_size, stride, dilation, is_causal
+        )
+        if embed_dim % num_heads != 0:
+            raise ValueError(
+                "Number of attention heads must evenly divide embedding dimension, "
+                f"got {embed_dim=}, {num_heads=}."
+            )
+        self.na_dim = na_dim
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_dim = self.embed_dim // self.num_heads
+        self.scale = qk_scale or self.head_dim**-0.5
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.dilation = dilation
+        self.is_causal = is_causal
+        self.expected_input_tensor_rank = self.na_dim + 2  # batch, embedding dim
+        self.qkv = nn.Linear(self.embed_dim, self.embed_dim * 3, bias=qkv_bias)
+        self.proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+    def forward(self, x: Tensor) -> Tensor:
+        if x.dim() != self.expected_input_tensor_rank:
+            raise ValueError(
+                f"NeighborhoodAttention{self.na_dim}D expected a tensor with rank "
+                f"{self.expected_input_tensor_rank} ({self.na_dim} for token layout, 1 for batch, "
+                f"1 for embedding dimension), got {x.dim()=}."
+            )
+        B, *input_shape, C = x.shape
+        if C != self.embed_dim:
+            raise ValueError(
+                f"Expected embedding dimension {self.embed_dim}, got {C} ({x.shape=})."
+            )
+        # 3, batch, *input_shape, heads, head_dim
+        permutation = (
+            [self.na_dim + 1, 0]
+            + [x + 1 for x in range(self.na_dim)]
+            + [self.na_dim + 2, self.na_dim + 3]
+        )
+        qkv = (
+            self.qkv(x)
+            .reshape(B, *input_shape, 3, self.num_heads, self.head_dim)
+            .permute(*permutation)
+        )
+        q, k, v = qkv[0], qkv[1], qkv[2]
+        x = neighborhood_attention_generic(  # type: ignore[assignment]
+            q,
+            k,
+            v,
+            kernel_size=self.kernel_size,
+            stride=self.stride,
+            dilation=self.dilation,
+            is_causal=self.is_causal,
+            scale=self.scale,
+        )
+        x = x.reshape(B, *input_shape, C)
+        return self.proj_drop(self.proj(x))
+    def extra_repr(self) -> str:
+        return (
+            f"head_dim={self.head_dim}, num_heads={self.num_heads}, "
+            + f"kernel_size={self.kernel_size}, "
+            + f"stride={self.stride}, "
+            + f"dilation={self.dilation}, "
+            + f"is_causal={self.is_causal}"
+        )
+class NeighborhoodAttention1D(NeighborhoodAttentionGeneric):
+    """
+    1-D Neighborhood Attention torch module.
+    Performs QKV and output linear projections in addition to the [na1d][natten.na1d] operation.
+    Args:
+        embed_dim: Embedding dimension size (a.k.a. number of channels, latent size).
+            !!! note
+                This is not `head_dim`. It's `head_dim * num_heads`.
+        num_heads: Number of attention heads.
+        kernel_size (Tuple[int] | int): Neighborhood window (kernel) size.
+            !!! note
+                `kernel_size` must be smaller than or equal to `seqlen`.
+        stride (Tuple[int] | int): Sliding window step size. Defaults to `1` (standard sliding
+            window).
+            !!! note
+                `stride` must be smaller than or equal to `kernel_size`.
+                When `stride == kernel_size`, there will be no overlap between sliding windows,
+                which is equivalent to blocked attention (a.k.a.
+                [window self attention](https://arxiv.org/abs/2103.14030)).
+        dilation (Tuple[int] | int): Dilation step size. Defaults to `1` (standard sliding window).
+            !!! note
+                The product of `dilation` and `kernel_size` must be smaller than or equal to
+                `seqlen`.
+        is_causal (Tuple[bool] | bool): Toggle causal masking. Defaults to `False`
+            (bi-directional).
+        qkv_bias: Enable bias in the QKV linear projection.
+        qk_scale: Attention scale. Defaults to `head_dim ** -0.5`.
+        proj_drop: Dropout score for projection layer. Defaults is `0.0` (no dropout).
+    Example:
+        ```python3
+        import torch
+        from . import NeighborhoodAttention1D
+        num_heads = 4
+        head_dim = 128
+        embed_dim = num_heads * head_dim
+        model = NeighborhoodAttention1D(
+            embed_dim=embed_dim,
+            num_heads=num_heads,
+            kernel_size=2048,
+            stride=2,
+            dilation=4,
+            is_causal=True
+        )
+        batch = 1
+        seqlen = 4096 # (1)!
+        x = torch.randn(batch, seqlen, embed_dim) # (2)!
+        y = model(x) # (3)!
+        ```
+        1. Tokens are arranged in a sequential layout of size 4096, to which we apply a
+            kernel size of 2048, stride 2, dilation 4, and apply causal masking.
+        2. `x.shape == [1, 4096, 512]`
+        3. `y.shape == [1, 4096, 512]`
+    """
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        kernel_size: Dimension1DTypeOrDed,
+        stride: Dimension1DTypeOrDed = 1,
+        dilation: Dimension1DTypeOrDed = 1,
+        is_causal: CausalArg1DTypeOrDed = False,
+        qkv_bias: bool = True,
+        qk_scale: Optional[float] = None,
+        proj_drop: float = 0.0,
+    ):
+        super().__init__(
+            na_dim=1,
+            embed_dim=embed_dim,
+            num_heads=num_heads,
+            kernel_size=kernel_size,
+            stride=stride,
+            dilation=dilation,
+            is_causal=is_causal,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            proj_drop=proj_drop,
+        )
+class NeighborhoodAttention2D(NeighborhoodAttentionGeneric):
+    """
+    2-D Neighborhood Attention torch module.
+    Performs QKV and output linear projections in addition to the [na2d][natten.na2d] operation.
+    Args:
+        embed_dim: Embedding dimension size (a.k.a. number of channels, latent size).
+            !!! note
+                This is not `head_dim`. It's `head_dim * num_heads`.
+        num_heads: Number of attention heads.
+        kernel_size (Tuple[int, int] | int): Neighborhood window (kernel) size/shape. If an
+            integer, it will be repeated for all 2 dimensions. For example `kernel_size=3` is
+            reinterpreted as `kernel_size=(3, 3)`.
+            !!! note
+                `kernel_size` must be smaller than or equal to token layout shape (`(X, Y)`) along
+                every dimension.
+        stride (Tuple[int, int] | int): Sliding window step size/shape. Defaults to `1` (standard
+            sliding window). If an integer, it will be repeated for all 2 dimensions. For example
+            `stride=2` is reinterpreted as `stride=(2, 2)`.
+            !!! note
+                `stride` must be smaller than or equal to `kernel_size` along every dimension.
+                When `stride == kernel_size`, there will be no overlap between sliding windows,
+                which is equivalent to blocked attention (a.k.a.
+                [window self attention](https://arxiv.org/abs/2103.14030)).
+        dilation (Tuple[int, int] | int): Dilation step size/shape. Defaults to `1` (standard
+            sliding window). If an integer, it will be repeated for all 2 dimensions. For example
+            `dilation=4` is reinterpreted as `dilation=(4, 4)`.
+            !!! note
+                The product of `dilation` and `kernel_size` must be smaller than or equal to
+                token layout shape (`(X, Y)`) along every dimension.
+        is_causal (Tuple[bool, bool] | bool): Toggle causal masking. Defaults to `False`
+            (bi-directional). If a boolean, it will be repeated for all 2 dimensions. For example
+            `is_causal=True` is reinterpreted as `is_causal=(True, True)`.
+        qkv_bias: Enable bias in the QKV linear projection.
+        qk_scale: Attention scale. Defaults to `head_dim ** -0.5`.
+        proj_drop: Dropout score for projection layer. Defaults is `0.0` (no dropout).
+    Example:
+        ```python3
+        import torch
+        from . import NeighborhoodAttention2D
+        num_heads = 4
+        head_dim = 128
+        embed_dim = num_heads * head_dim
+        model = NeighborhoodAttention2D(
+            embed_dim=embed_dim,
+            num_heads=num_heads,
+            kernel_size=(8, 16),
+            stride=(1, 2),
+            dilation=(2, 1),
+            is_causal=False
+        )
+        batch = 1
+        token_layout_shape = (16, 32) # (1)!
+        x = torch.randn(batch, *token_layout_shape, embed_dim) # (2)!
+        y = model(x) # (3)!
+        ```
+        1. Tokens are arranged in a 16 x 32 layout, to which we apply a
+            kernel size of 8 x 16,
+            stride 1 x 2,
+            and dilation 2 x 1.
+        2. `x.shape == [1, 16, 32, 512]`
+        3. `y.shape == [1, 16, 32, 512]`
+    """
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        kernel_size: Dimension2DTypeOrDed,
+        stride: Dimension2DTypeOrDed = 1,
+        dilation: Dimension2DTypeOrDed = 1,
+        is_causal: CausalArg2DTypeOrDed = False,
+        qkv_bias: bool = True,
+        qk_scale: Optional[float] = None,
+        proj_drop: float = 0.0,
+    ):
+        super().__init__(
+            na_dim=2,
+            embed_dim=embed_dim,
+            num_heads=num_heads,
+            kernel_size=kernel_size,
+            stride=stride,
+            dilation=dilation,
+            is_causal=is_causal,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            proj_drop=proj_drop,
+        )
+class NeighborhoodAttention3D(NeighborhoodAttentionGeneric):
+    """
+    3-D Neighborhood Attention torch module.
+    Performs QKV and output linear projections in addition to the [na3d][natten.na3d] operation.
+    Args:
+        embed_dim: Embedding dimension size (a.k.a. number of channels, latent size).
+            !!! note
+                This is not `head_dim`. It's `head_dim * num_heads`.
+        num_heads: Number of attention heads.
+        kernel_size (Tuple[int, int, int] | int): Neighborhood window (kernel) size/shape. If an
+            integer, it will be repeated for all 3 dimensions. For example `kernel_size=3` is
+            reinterpreted as `kernel_size=(3, 3, 3)`.
+            !!! note
+                `kernel_size` must be smaller than or equal to token layout shape (`(X, Y, Z)`)
+                along every dimension.
+        stride (Tuple[int, int, int] | int): Sliding window step size/shape. Defaults to `1`
+            (standard sliding window). If an integer, it will be repeated for all 3 dimensions.
+            For example `stride=2` is reinterpreted as `stride=(2, 2, 2)`.
+            !!! note
+                `stride` must be smaller than or equal to `kernel_size` along every dimension.
+                When `stride == kernel_size`, there will be no overlap between sliding windows,
+                which is equivalent to blocked attention (a.k.a.
+                [window self attention](https://arxiv.org/abs/2103.14030)).
+        dilation (Tuple[int, int, int] | int): Dilation step size/shape. Defaults to `1` (standard
+            sliding window). If an integer, it will be repeated for all 3 dimensions. For example
+            `dilation=4` is reinterpreted as `dilation=(4, 4, 4)`.
+            !!! note
+                The product of `dilation` and `kernel_size` must be smaller than or equal to
+                token layout shape (`(X, Y, Z)`) along every dimension.
+        is_causal (Tuple[bool, bool, bool] | bool): Toggle causal masking. Defaults to `False`
+            (bi-directional). If a boolean, it will be repeated for all 3 dimensions. For example
+            `is_causal=True` is reinterpreted as `is_causal=(True, True, True)`.
+        qkv_bias: Enable bias in the QKV linear projection.
+        qk_scale: Attention scale. Defaults to `head_dim ** -0.5`.
+        proj_drop: Dropout score for projection layer. Defaults is `0.0` (no dropout).
+    Example:
+        ```python3
+        import torch
+        from . import NeighborhoodAttention3D
+        num_heads = 4
+        head_dim = 128
+        embed_dim = num_heads * head_dim
+        model = NeighborhoodAttention3D(
+            embed_dim=embed_dim,
+            num_heads=num_heads,
+            kernel_size=(4, 8, 12),
+            stride=(1, 1, 4),
+            dilation=(1, 2, 1),
+            is_causal=(True, False, False)
+        )
+        batch = 1
+        token_layout_shape = (12, 16, 20) # (1)!
+        x = torch.randn(batch, *token_layout_shape, embed_dim) # (2)!
+        y = model(x) # (3)!
+        ```
+        1. Tokens are arranged in a 12 x 16 x 20 layout, to which we apply a
+            kernel size of 4 x 8 x 12,
+            stride 1 x 1 x 4,
+            dilation 1 x 2 x 1, and apply causal masking to the left-most dimension (12).
+        2. `x.shape == [1, 12, 16, 20, 512]`
+        3. `y.shape == [1, 12, 16, 20, 512]`
+    """
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        kernel_size: Dimension3DTypeOrDed,
+        stride: Dimension3DTypeOrDed = 1,
+        dilation: Dimension3DTypeOrDed = 1,
+        is_causal: CausalArg3DTypeOrDed = False,
+        qkv_bias: bool = True,
+        qk_scale: Optional[float] = None,
+        proj_drop: float = 0.0,
+    ):
+        super().__init__(
+            na_dim=3,
+            embed_dim=embed_dim,
+            num_heads=num_heads,
+            kernel_size=kernel_size,
+            stride=stride,
+            dilation=dilation,
+            is_causal=is_causal,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            proj_drop=proj_drop,
+        )

build/torch210-cu126-aarch64-linux/symbols.json ADDED Viewed

	@@ -0,0 +1,1257 @@

+{
+  "schema_version": 1,
+  "module": "_natten_cuda_5fe9dbb",
+  "functions": [
+    {
+      "name": "are_deterministic_algorithms_enabled",
+      "qualname": "are_deterministic_algorithms_enabled",
+      "module": "_natten_cuda_5fe9dbb.context",
+      "kind": "function",
+      "doc": null,
+      "signature": {
+        "parameters": [],
+        "return_annotation": "bool"
+      }
+    },
+    {
+      "name": "use_deterministic_algorithms",
+      "qualname": "use_deterministic_algorithms",
+      "module": "_natten_cuda_5fe9dbb.context",
+      "kind": "function",
+      "doc": null,
+      "signature": {
+        "parameters": [
+          {
+            "name": "mode",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "True",
+            "annotation": "bool"
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "use_kv_parallelism_in_fused_na",
+      "qualname": "use_kv_parallelism_in_fused_na",
+      "module": "_natten_cuda_5fe9dbb.context",
+      "kind": "function",
+      "doc": "Sets guards for using KV Parallelism in backpropagation in `\"cutlass-fna\"`/`\"cutlass-fmha\"`\nbackends.\n\nWarning:\n    Disabling KV parallelism can significantly slow down training, particularly in\n    small-batch/head and large-token problems.\n\nArgs:\n    mode: If `True`, allows KV parallelism (default setting), and otherwise disables it.",
+      "signature": {
+        "parameters": [
+          {
+            "name": "mode",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "True",
+            "annotation": "bool"
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "is_kv_parallelism_in_fused_na_enabled",
+      "qualname": "is_kv_parallelism_in_fused_na_enabled",
+      "module": "_natten_cuda_5fe9dbb.context",
+      "kind": "function",
+      "doc": "Returns whether KV parallelism in `\"cutlass-fna\"` and `\"cutlass-fmha\"` backends is enabled.",
+      "signature": {
+        "parameters": [],
+        "return_annotation": "bool"
+      }
+    },
+    {
+      "name": "set_memory_usage_preference",
+      "qualname": "set_memory_usage_preference",
+      "module": "_natten_cuda_5fe9dbb.context",
+      "kind": "function",
+      "doc": "Sets memory usage preference for KV parallelism in `\"cutlass-fna\"` and `\"cutlass-fmha\"`\nbackends.\n\nArgs:\n    pref: Choices are `\"default\"`, `\"strict\"`, and `\"unrestricted\"`.",
+      "signature": {
+        "parameters": [
+          {
+            "name": "pref",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "'default'",
+            "annotation": "str"
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "get_memory_usage_preference",
+      "qualname": "get_memory_usage_preference",
+      "module": "_natten_cuda_5fe9dbb.context",
+      "kind": "function",
+      "doc": null,
+      "signature": {
+        "parameters": [],
+        "return_annotation": "_natten_cuda_5fe9dbb.context.MemoryUsagePreference"
+      }
+    },
+    {
+      "name": "is_memory_usage_default",
+      "qualname": "is_memory_usage_default",
+      "module": "_natten_cuda_5fe9dbb.context",
+      "kind": "function",
+      "doc": "Returns whether memory usage preference for KV parallelism in `\"cutlass-fna\"` and\n`\"cutlass-fmha\"` backends is the default setting.",
+      "signature": {
+        "parameters": [],
+        "return_annotation": "bool"
+      }
+    },
+    {
+      "name": "is_memory_usage_strict",
+      "qualname": "is_memory_usage_strict",
+      "module": "_natten_cuda_5fe9dbb.context",
+      "kind": "function",
+      "doc": "Returns whether memory usage preference for KV parallelism in `\"cutlass-fna\"` and\n`\"cutlass-fmha\"` backends is the *restricted* setting.",
+      "signature": {
+        "parameters": [],
+        "return_annotation": "bool"
+      }
+    },
+    {
+      "name": "is_memory_usage_unrestricted",
+      "qualname": "is_memory_usage_unrestricted",
+      "module": "_natten_cuda_5fe9dbb.context",
+      "kind": "function",
+      "doc": "Returns whether memory usage preference for KV parallelism in `\"cutlass-fna\"` and\n`\"cutlass-fmha\"` backends is the *unrestricted* setting.",
+      "signature": {
+        "parameters": [],
+        "return_annotation": "bool"
+      }
+    },
+    {
+      "name": "is_flex_compile_allowed",
+      "qualname": "is_flex_compile_allowed",
+      "module": "_natten_cuda_5fe9dbb.context",
+      "kind": "function",
+      "doc": "Returns whether compilation is allowed in `\"flex-fna\"` and `\"flex-fmha\"` backends.",
+      "signature": {
+        "parameters": [],
+        "return_annotation": "bool"
+      }
+    },
+    {
+      "name": "is_flex_compile_backprop_allowed",
+      "qualname": "is_flex_compile_backprop_allowed",
+      "module": "_natten_cuda_5fe9dbb.context",
+      "kind": "function",
+      "doc": "Returns whether compilation for backpropagation is allowed in `\"flex-fna\"` and `\"flex-fmha\"`\nbackends.",
+      "signature": {
+        "parameters": [],
+        "return_annotation": "bool"
+      }
+    },
+    {
+      "name": "allow_flex_compile",
+      "qualname": "allow_flex_compile",
+      "module": "_natten_cuda_5fe9dbb.context",
+      "kind": "function",
+      "doc": "Sets guards for Flex Attention + `torch.compile`.\n\nAllows using our Flex FNA / Flex FMHA backends with `torch.compile`, meaning you can\npass `torch_compile=True` to the `na{1,2,3}d` or `attention` operation, along with\n`backend=\"flex-fna\"`/`backend=\"flex-fmha\"`, and NATTEN will compile the block-sparse mask, as\nwell as the attention operation using `torch.compile` for you.\n\nWarning:\n    We have been *unable to verify the correctness* of this setting under all of our use\n    cases. We are working on raising this issue with PyTorch directly, but until then we strongly\n    recommend exercising caution when using this feature.\n\nDanger: backprop=True is strongly discouraged!\n    Allowing `torch.compile` for backpropagation (detected by checking\n    `tensor.requires_grad`) is guarded separately. We strongly recommend NOT using this setting, as\n    it can impact your training results.\n\nArgs:\n    mode: If `True`, enable compilation for forward pass, otherwise disable.\n    backprop: If `True`, assuming compilation for forward pass is allowed, enable compilation\n        for backward pass, otherwise disable.",
+      "signature": {
+        "parameters": [
+          {
+            "name": "mode",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "True",
+            "annotation": "bool"
+          },
+          {
+            "name": "backprop",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "allow_flex_compile_backprop",
+      "qualname": "allow_flex_compile_backprop",
+      "module": "_natten_cuda_5fe9dbb.context",
+      "kind": "function",
+      "doc": "Sets guards for Flex Attention + `torch.compile` for backpropagation only.\n\nArgs:\n    mode: If `True`, enable compilation for backprop (assuming forward compilation is already\n        enabled), otherwise disable.",
+      "signature": {
+        "parameters": [
+          {
+            "name": "mode",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "True",
+            "annotation": "bool"
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "disable_flex_compile",
+      "qualname": "disable_flex_compile",
+      "module": "_natten_cuda_5fe9dbb.context",
+      "kind": "function",
+      "doc": "Disallow Flex Attention + `torch.compile` entirely.",
+      "signature": {
+        "parameters": [],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "disable_flex_compile_backprop",
+      "qualname": "disable_flex_compile_backprop",
+      "module": "_natten_cuda_5fe9dbb.context",
+      "kind": "function",
+      "doc": "Disallow Flex Attention + `torch.compile` for backpropagation entirely.",
+      "signature": {
+        "parameters": [],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "get_bwd_configs_for_cutlass_fmha",
+      "qualname": "get_bwd_configs_for_cutlass_fmha",
+      "module": "_natten_cuda_5fe9dbb.backends.configs",
+      "kind": "function",
+      "doc": "Returns CUTLASS FMHA backward pass configurations compatible with input tensors, if any.\n\nChecks first if a CUDA tensor, and on a device with compute capability >= 5.0, and if so,\nreturns *backward pass* configurations compatible with the specific compute capability, tensor\ndtype and head dim.\n\nEach configuration for this operation is a tuple of two integers: `(backward_q_tile_size,\nbackward_kv_tile_size)`. These are arguments to [natten.attention][natten.attention].\n\nArgs:\n    query: Query tensor matching the shape, dtype, and device of your use case.\n    key:   Key tensor matching the shape, dtype, and device of your use case.\n    value: Value tensor matching the shape, dtype, and device of your use case.\n\nReturns:\n    (List[Tuple[int, int]]): List of tuples of two integers corresponding to query and KV tile\n        sizes in the *backward pass*.",
+      "signature": {
+        "parameters": [
+          {
+            "name": "query",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "key",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "value",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          }
+        ],
+        "return_annotation": "List[Tuple[int, int]]"
+      }
+    },
+    {
+      "name": "get_bwd_configs_for_cutlass_fna",
+      "qualname": "get_bwd_configs_for_cutlass_fna",
+      "module": "_natten_cuda_5fe9dbb.backends.configs",
+      "kind": "function",
+      "doc": "Returns CUTLASS FNA backward pass configurations compatible with input tensors, if any.\n\nChecks first if a CUDA tensor, and on a device with compute capability >= 5.0, and if so,\nreturns *backward pass* configurations compatible with the specific compute capability, tensor\ndtype and head dim.\n\nEach configuration for this operation is a tuple of two integers: `(backward_q_tile_shape,\nbackward_kv_tile_shape)`. These are arguments to [natten.na1d][natten.na1d],\n[natten.na2d][natten.na2d], and [natten.na3d][natten.na3d].\n\nArgs:\n    query: Query tensor matching the shape, dtype, and device of your use case.\n    key:   Key tensor matching the shape, dtype, and device of your use case.\n    value: Value tensor matching the shape, dtype, and device of your use case.\n\nReturns:\n    (List[Tuple[tuple, tuple]]): List of tuples of two integer tuples corresponding to query\n        and KV tile *shapes* in the *backward pass*.",
+      "signature": {
+        "parameters": [
+          {
+            "name": "query",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "key",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "value",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          }
+        ],
+        "return_annotation": "List[Tuple[Tuple[int], Tuple[int]] | Tuple[Tuple[int, int], Tuple[int, int]] | Tuple[Tuple[int, int, int], Tuple[int, int, int]]]"
+      }
+    },
+    {
+      "name": "get_configs_for_cutlass_fmha",
+      "qualname": "get_configs_for_cutlass_fmha",
+      "module": "_natten_cuda_5fe9dbb.backends.configs",
+      "kind": "function",
+      "doc": "Returns CUTLASS FMHA configurations compatible with input tensors, if any.\n\nChecks first if a CUDA tensor, and on a device with compute capability >= 5.0, and if so,\nreturns *forward pass* configurations compatible with the specific compute capability, tensor\ndtype and head dim.\n\nEach configuration for this operation is a tuple of two integers: `(q_tile_size,\nkv_tile_size)`. These are arguments to [natten.attention][natten.attention].\n\nArgs:\n    query: Query tensor matching the shape, dtype, and device of your use case.\n    key:   Key tensor matching the shape, dtype, and device of your use case.\n    value: Value tensor matching the shape, dtype, and device of your use case.\n\nReturns:\n    (List[Tuple[int, int]]): List of tuples of two integers corresponding to query and KV tile\n        sizes.",
+      "signature": {
+        "parameters": [
+          {
+            "name": "query",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "key",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "value",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          }
+        ],
+        "return_annotation": "List[Tuple[int, int]]"
+      }
+    },
+    {
+      "name": "get_configs_for_cutlass_fna",
+      "qualname": "get_configs_for_cutlass_fna",
+      "module": "_natten_cuda_5fe9dbb.backends.configs",
+      "kind": "function",
+      "doc": "Returns CUTLASS FNA configurations compatible with input tensors, if any.\n\nChecks first if a CUDA tensor, and on a device with compute capability >= 5.0, and if so,\nreturns *forward pass* configurations compatible with the specific compute capability, tensor\ndtype and head dim, and according to the rank of the token layout (1D/2D/3D).\n\nEach configuration for this operation is a tuple of two integer tuples: `(q_tile_shape,\nkv_tile_shape)`. These are arguments to [natten.na1d][natten.na1d], [natten.na2d][natten.na2d],\nand [natten.na3d][natten.na3d].\n\nArgs:\n    query: Query tensor matching the shape, dtype, and device of your use case.\n    key:   Key tensor matching the shape, dtype, and device of your use case.\n    value: Value tensor matching the shape, dtype, and device of your use case.\n\nReturns:\n    (List[Tuple[tuple, tuple]]): List of tuples of two integer tuples corresponding to query\n        and KV tile *shapes*.",
+      "signature": {
+        "parameters": [
+          {
+            "name": "query",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "key",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "value",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          }
+        ],
+        "return_annotation": "List[Tuple[Tuple[int], Tuple[int]] | Tuple[Tuple[int, int], Tuple[int, int]] | Tuple[Tuple[int, int, int], Tuple[int, int, int]]]"
+      }
+    },
+    {
+      "name": "get_configs_for_cutlass_hopper_fmha",
+      "qualname": "get_configs_for_cutlass_hopper_fmha",
+      "module": "_natten_cuda_5fe9dbb.backends.configs",
+      "kind": "function",
+      "doc": "Returns Hopper FMHA configurations compatible with input tensors, if any.\n\nChecks first if a CUDA tensor, and on a Hopper GPU (SM90; compute capability 9.0), and if so,\nreturns *forward pass* configurations compatible with the tensor dtype and head dim.\n\nEach configuration for this operation is a tuple of one integer tuple, and another integer:\n`((q_tile_size, kv_tile_size), kernel_schedule)`. These are arguments to\n[natten.attention][natten.attention].\n`kernel_schedule` is specific to Hopper FNA/FMHA only.\n\nArgs:\n    query: Query tensor matching the shape, dtype, and device of your use case.\n    key:   Key tensor matching the shape, dtype, and device of your use case.\n    value: Value tensor matching the shape, dtype, and device of your use case.\n\nReturns:\n    (List[Tuple[Tuple[int, int], KernelSchedule]]): List of tuples of one tuple of two integers\n        corresponding to query and KV tile sizes, and a kernel schedule enum type.",
+      "signature": {
+        "parameters": [
+          {
+            "name": "query",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "key",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "value",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          }
+        ],
+        "return_annotation": "List[Tuple[Tuple[int, int], _natten_cuda_5fe9dbb._types.KernelSchedule]]"
+      }
+    },
+    {
+      "name": "get_bwd_configs_for_cutlass_hopper_fmha",
+      "qualname": "get_bwd_configs_for_cutlass_hopper_fmha",
+      "module": "_natten_cuda_5fe9dbb.backends.configs",
+      "kind": "function",
+      "doc": "Returns Hopper FMHA backward pass configurations compatible with input tensors, if any.\n\nChecks first if a CUDA tensor, and on a Hopper GPU (SM90; compute capability 9.0), and if so,\nreturns *backward pass* configurations compatible with the tensor dtype and head dim.\n\nEach configuration for this operation is an integer tuple:\n`(backward_q_tile_size, backward_kv_tile_size)`. These are arguments to\n[natten.attention][natten.attention].\n\nNote that unlike forward pass, kernel schedule is not part of the configuration. All backward\npass kernels are persistent warp-specialized.\nSee CUTLASS's [example 88](https://github.com/NVIDIA/cutlass/tree/main/examples/88_hopper_fmha)\nfor more.\n\nArgs:\n    query: Query tensor matching the shape, dtype, and device of your use case.\n    key:   Key tensor matching the shape, dtype, and device of your use case.\n    value: Value tensor matching the shape, dtype, and device of your use case.\n\nReturns:\n    (List[Tuple[int, int]]): List of integer tuples corresponding to query and KV tile sizes.",
+      "signature": {
+        "parameters": [
+          {
+            "name": "query",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "key",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "value",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          }
+        ],
+        "return_annotation": "List[Tuple[int, int]]"
+      }
+    },
+    {
+      "name": "get_configs_for_cutlass_hopper_fna",
+      "qualname": "get_configs_for_cutlass_hopper_fna",
+      "module": "_natten_cuda_5fe9dbb.backends.configs",
+      "kind": "function",
+      "doc": "Returns Hopper FNA configurations compatible with input tensors, if any.\n\nChecks first if a CUDA tensor, and on a Hopper GPU (SM90; compute capability 9.0), and if so,\nreturns *forward pass* configurations compatible with the tensor dtype and head dim.\n\nEach configuration for this operation is a tuple of one tuple, and another integer:\n`((q_tile_shape, kv_tile_shape), kernel_schedule)`. These are arguments to\n[natten.na1d][natten.na1d], [natten.na2d][natten.na2d], and [natten.na3d][natten.na3d].\n\n`kernel_schedule` is specific to Hopper FNA/FMHA only.\n\nArgs:\n    query: Query tensor matching the shape, dtype, and device of your use case.\n    key:   Key tensor matching the shape, dtype, and device of your use case.\n    value: Value tensor matching the shape, dtype, and device of your use case.\n\nReturns:\n    (List[Tuple[Tuple[tuple, tuple], KernelSchedule]]): List of tuples of one tuple of two\n        shape tuples, corresponding to query and KV tile *shapes*, and a kernel schedule enum\n        type.",
+      "signature": {
+        "parameters": [
+          {
+            "name": "query",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "key",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "value",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          }
+        ],
+        "return_annotation": "List[Tuple[Tuple[Tuple[int], Tuple[int]] | Tuple[Tuple[int, int], Tuple[int, int]] | Tuple[Tuple[int, int, int], Tuple[int, int, int]], _natten_cuda_5fe9dbb._types.KernelSchedule]]"
+      }
+    },
+    {
+      "name": "get_bwd_configs_for_cutlass_hopper_fna",
+      "qualname": "get_bwd_configs_for_cutlass_hopper_fna",
+      "module": "_natten_cuda_5fe9dbb.backends.configs",
+      "kind": "function",
+      "doc": "Returns Hopper FNA backward pass configurations compatible with input tensors, if any.\n\nChecks first if a CUDA tensor, and on a Hopper GPU (SM90; compute capability 9.0), and if so,\nreturns *backward pass* configurations compatible with the tensor dtype and head dim.\n\nEach configuration for this operation is a tuple of two tuples:\n`(q_tile_shape, kv_tile_shape)`. These are arguments to [natten.na1d][natten.na1d],\n[natten.na2d][natten.na2d], and [natten.na3d][natten.na3d].\n\nNote that unlike forward pass, kernel schedule is not part of the configuration. All backward\npass kernels are persistent warp-specialized.\nSee CUTLASS's [example 88](https://github.com/NVIDIA/cutlass/tree/main/examples/88_hopper_fmha)\nfor more.\n\nArgs:\n    query: Query tensor matching the shape, dtype, and device of your use case.\n    key:   Key tensor matching the shape, dtype, and device of your use case.\n    value: Value tensor matching the shape, dtype, and device of your use case.\n\nReturns:\n    (List[Tuple[tuple, tuple]]): List of tuples of two shape tuples, corresponding to query and\n    KV tile *shapes*.",
+      "signature": {
+        "parameters": [
+          {
+            "name": "query",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "key",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "value",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          }
+        ],
+        "return_annotation": "List[Tuple[Tuple[int], Tuple[int]] | Tuple[Tuple[int, int], Tuple[int, int]] | Tuple[Tuple[int, int, int], Tuple[int, int, int]]]"
+      }
+    },
+    {
+      "name": "get_bwd_configs_for_cutlass_blackwell_fmha",
+      "qualname": "get_bwd_configs_for_cutlass_blackwell_fmha",
+      "module": "_natten_cuda_5fe9dbb.backends.configs",
+      "kind": "function",
+      "doc": "Returns Blackwell FMHA backward pass configurations compatible with input tensors, if any.\n\nChecks first if a CUDA tensor, and on a Blackwell datacenter GPU (SM100; compute capability\n10.0), and if so, returns *backward pass* configurations compatible with the tensor dtype and\nhead dim.\n\nEach configuration for this operation is a tuple of two integers: `(q_tile_size,\nkv_tile_size)`. These are arguments to [natten.attention][natten.attention].\n\nArgs:\n    query: Query tensor matching the shape, dtype, and device of your use case.\n    key:   Key tensor matching the shape, dtype, and device of your use case.\n    value: Value tensor matching the shape, dtype, and device of your use case.\n\nReturns:\n    (List[Tuple[int, int]]): List of tuples of two integers corresponding to query and KV tile\n        sizes.",
+      "signature": {
+        "parameters": [
+          {
+            "name": "query",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "key",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "value",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          }
+        ],
+        "return_annotation": "List[Tuple[int, int]]"
+      }
+    },
+    {
+      "name": "get_bwd_configs_for_cutlass_blackwell_fna",
+      "qualname": "get_bwd_configs_for_cutlass_blackwell_fna",
+      "module": "_natten_cuda_5fe9dbb.backends.configs",
+      "kind": "function",
+      "doc": "Returns Blackwell FNA backward pass configurations compatible with input tensors, if any.\n\nChecks first if a CUDA tensor, and on a Blackwell datacenter GPU (SM100; compute capability\n10.0), and if so, returns *backward pass* configurations compatible with the tensor dtype and\nhead dim, and according to the rank of the token layout (1D/2D/3D).\n\nEach configuration for this operation is a tuple of two integer tuples: `(q_tile_shape,\nkv_tile_shape)`. These are arguments to [natten.na1d][natten.na1d], [natten.na2d][natten.na2d],\nand [natten.na3d][natten.na3d].\n\nArgs:\n    query: Query tensor matching the shape, dtype, and device of your use case.\n    key:   Key tensor matching the shape, dtype, and device of your use case.\n    value: Value tensor matching the shape, dtype, and device of your use case.\n\nReturns:\n    (List[Tuple[tuple, tuple]]): List of tuples of two integer tuples corresponding to query\n        and KV tile *shapes*.",
+      "signature": {
+        "parameters": [
+          {
+            "name": "query",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "key",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "value",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          }
+        ],
+        "return_annotation": "List[Tuple[Tuple[int], Tuple[int]] | Tuple[Tuple[int, int], Tuple[int, int]] | Tuple[Tuple[int, int, int], Tuple[int, int, int]]]"
+      }
+    },
+    {
+      "name": "get_configs_for_cutlass_blackwell_fmha",
+      "qualname": "get_configs_for_cutlass_blackwell_fmha",
+      "module": "_natten_cuda_5fe9dbb.backends.configs",
+      "kind": "function",
+      "doc": "Returns Blackwell FMHA configurations compatible with input tensors, if any.\n\nChecks first if a CUDA tensor, and on a Blackwell datacenter GPU (SM100; compute capability\n10.0), and if so, returns *forward pass* configurations compatible with the tensor dtype and\nhead dim.\n\nEach configuration for this operation is a tuple of two integers: `(q_tile_size,\nkv_tile_size)`. These are arguments to [natten.attention][natten.attention].\n\nArgs:\n    query: Query tensor matching the shape, dtype, and device of your use case.\n    key:   Key tensor matching the shape, dtype, and device of your use case.\n    value: Value tensor matching the shape, dtype, and device of your use case.\n\nReturns:\n    (List[Tuple[int, int]]): List of tuples of two integers corresponding to query and KV tile\n        sizes.",
+      "signature": {
+        "parameters": [
+          {
+            "name": "query",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "key",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "value",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          }
+        ],
+        "return_annotation": "List[Tuple[int, int]]"
+      }
+    },
+    {
+      "name": "get_configs_for_cutlass_blackwell_fna",
+      "qualname": "get_configs_for_cutlass_blackwell_fna",
+      "module": "_natten_cuda_5fe9dbb.backends.configs",
+      "kind": "function",
+      "doc": "Returns Blackwell FNA configurations compatible with input tensors, if any.\n\nChecks first if a CUDA tensor, and on a Blackwell datacenter GPU (SM100; compute capability\n10.0), and if so, returns *forward pass* configurations compatible with the tensor dtype and\nhead dim, and according to the rank of the token layout (1D/2D/3D).\n\nEach configuration for this operation is a tuple of two integer tuples: `(q_tile_shape,\nkv_tile_shape)`. These are arguments to [natten.na1d][natten.na1d], [natten.na2d][natten.na2d],\nand [natten.na3d][natten.na3d].\n\nArgs:\n    query: Query tensor matching the shape, dtype, and device of your use case.\n    key:   Key tensor matching the shape, dtype, and device of your use case.\n    value: Value tensor matching the shape, dtype, and device of your use case.\n\nReturns:\n    (List[Tuple[tuple, tuple]]): List of tuples of two integer tuples corresponding to query\n        and KV tile *shapes*.",
+      "signature": {
+        "parameters": [
+          {
+            "name": "query",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "key",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "value",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          }
+        ],
+        "return_annotation": "List[Tuple[Tuple[int], Tuple[int]] | Tuple[Tuple[int, int], Tuple[int, int]] | Tuple[Tuple[int, int, int], Tuple[int, int, int]]]"
+      }
+    },
+    {
+      "name": "get_configs_for_flex_fmha",
+      "qualname": "get_configs_for_flex_fmha",
+      "module": "_natten_cuda_5fe9dbb.backends.configs",
+      "kind": "function",
+      "doc": "Returns Flex FMHA configurations compatible with input tensors, if any.\n\nEach configuration for this operation is a tuple of two integers: `(q_tile_size,\nkv_tile_size)`. These are arguments to [natten.attention][natten.attention].\nNot specifying these arguments while backend is Flex will default to `q_tile_size = 64` and\n`kv_tile_size = 64`.\n\nArgs:\n    query: Query tensor matching the shape, dtype, and device of your use case.\n    key:   Key tensor matching the shape, dtype, and device of your use case.\n    value: Value tensor matching the shape, dtype, and device of your use case.\n    torch_compile: Whether or not you intend to use compiled block mask and flex attention kernel.\n\nReturns:\n    (List[Tuple[int, int]]): List of tuples of two integers corresponding to query and KV tile\n        sizes.",
+      "signature": {
+        "parameters": [
+          {
+            "name": "query",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "key",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "value",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "torch_compile",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          }
+        ],
+        "return_annotation": "List[Tuple[int, int]]"
+      }
+    },
+    {
+      "name": "get_configs_for_flex_fna",
+      "qualname": "get_configs_for_flex_fna",
+      "module": "_natten_cuda_5fe9dbb.backends.configs",
+      "kind": "function",
+      "doc": "Returns Flex FNA configurations compatible with input tensors, if any.\n\nEach configuration for this operation is a tuple of two integer tuples: `(q_tile_shape,\nkv_tile_shape)`. These are arguments to [natten.na1d][natten.na1d], [natten.na2d][natten.na2d],\nand [natten.na3d][natten.na3d].\nNot specifying these arguments while backend is Flex will default to single-dimensional tiling,\nand will not use our Token Permutation approach. By explicitly specifying tile shapes, you will\nautomatically use our Token Permutation approach, which saves you the most compute.\n\nArgs:\n    query: Query tensor matching the shape, dtype, and device of your use case.\n    key:   Key tensor matching the shape, dtype, and device of your use case.\n    value: Value tensor matching the shape, dtype, and device of your use case.\n    torch_compile: Whether or not you intend to use compiled block mask and flex attention kernel.\n\nReturns:\n    (List[Tuple[tuple, tuple]]): List of tuples of two integer tuples corresponding to query\n        and KV tile *shapes*.",
+      "signature": {
+        "parameters": [
+          {
+            "name": "query",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "key",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "value",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "torch_compile",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          }
+        ],
+        "return_annotation": "List[Tuple[Tuple[int], Tuple[int]] | Tuple[Tuple[int, int], Tuple[int, int]] | Tuple[Tuple[int, int, int], Tuple[int, int, int]]]"
+      }
+    },
+    {
+      "name": "na1d",
+      "qualname": "na1d",
+      "module": "_natten_cuda_5fe9dbb.functional",
+      "kind": "function",
+      "doc": "Computes 1-D neighborhood attention.\n\nGQA/MQA support (`heads != heads_kv`) is available. For now, `blackwell-fna` and\n`flex-fna` support GQA/MQA natively, and `cutlass-fna` and `hopper-fna` support it with\nexplicit repeats (increases memory usage and runtime).\n\nParameters:\n    query (Tensor): 4-D query tensor, with the heads last layout\n        (`[batch, seqlen, heads, head_dim]`)\n\n    key (Tensor): 4-D key tensor, with the heads last layout\n        (`[batch, seqlen, heads_kv, head_dim]`)\n\n    value (Tensor): 4-D value tensor, with the heads last layout\n        (`[batch, seqlen, heads_kv, head_dim_v]`)\n\n    kernel_size (Tuple[int] | int): Neighborhood window (kernel) size.\n\n        !!! note\n            `kernel_size` must be smaller than or equal to `seqlen`.\n\n    stride (Tuple[int] | int): Sliding window step size. Defaults to `1` (standard sliding\n        window).\n\n        !!! note\n            `stride` must be smaller than or equal to `kernel_size`.\n            When `stride == kernel_size`, there will be no overlap between sliding windows,\n            which is equivalent to blocked attention (a.k.a.\n            [window self attention](https://arxiv.org/abs/2103.14030)).\n\n    dilation (Tuple[int] | int): Dilation step size. Defaults to `1` (standard sliding window).\n\n        !!! note\n            The product of `dilation` and `kernel_size` must be smaller than or equal to\n            `seqlen`.\n\n    is_causal (Tuple[bool] | bool): Toggle causal masking. Defaults to `False`\n        (bi-directional).\n\n    scale (float): Attention scale. Defaults to `head_dim ** -0.5`.\n\n    additional_keys: `None` or 4-D key tensor, with the heads last layout\n        (`[batch, seqlen_kv, heads_kv, head_dim]`), corresponding to key tokens from some\n        additional context. Used when performing neighborhood cross-attention, where `query`\n        tokens attend to their neighborhood, as well as some fixed additional set of tokens.\n\n    additional_values: `None` or 4-D value tensor, with the heads last layout\n        (`[batch, seqlen_kv, heads_kv, head_dim_v]`), corresponding to value tokens from some\n        additional context. Used when performing neighborhood cross-attention, where `query`\n        tokens attend to their neighborhood, as well as some fixed additional set of tokens.\n\n        !!! note\n            `additional_keys` and `additional_values` must both either be `Tensor`s, or both\n            `None`s, and must match in shape.\n\nOther Parameters:\n    backend (str): Backend implementation to run with. Choices are: `None` (pick the best\n        available one), `\"cutlass-fna\"`, `\"hopper-fna\"`, `\"blackwell-fna\"`, `\"flex-fna\"`.\n        Refer to [backends](backends.md) for more information.\n\n    q_tile_shape (Tuple[int]): 1-D Tile shape for the query token layout in the forward pass\n        kernel. You can use [profiler](profiler.md) to find valid choices for your use case,\n        and search for the best combination.\n\n    kv_tile_shape (Tuple[int]): 1-D Tile shape for the key-value token layout in the forward\n        pass kernel. You can use [profiler](profiler.md) to find valid choices for your use\n        case, and search for the best combination.\n\n    backward_q_tile_shape (Tuple[int]): 1-D Tile shape for the query token layout in the\n        backward pass kernel. This is ignored by `\"flex-fna\"`.\n        You can use [profiler](profiler.md) to find valid choices for your use case, and\n        search for the best combination.\n\n    backward_kv_tile_shape (Tuple[int]): 1-D Tile shape for the key/value token layout in the\n        backward pass kernel. This is ignored by `\"flex-fna\"`.\n        You can use [profiler](profiler.md) to find valid choices for your use case, and\n        search for the best combination.\n\n    backward_kv_splits (Tuple[int]): Number of key/value tiles allowed to work in parallel in\n        the backward pass kernel. Like tile shapes, this is a tuple and not an integer for\n        neighborhood attention operations, and the size of the tuple corresponds to the number\n        of dimensions / rank of the layout of tokens. This is only respected by the\n        `\"cutlass-fna\"` backend, and only when\n        [KV parallelism](context.md#kv-parallelism-in-fna) is enabled.\n\n    backward_use_pt_reduction (bool): Whether to use PyTorch eager for computing the `dO * O`\n        product required by the backward pass, over the CUTLASS kernel. This only applies to\n        the `\"cutlass-fna\"` backend.\n\n    run_persistent_kernel (bool): Whether to use persistent tile scheduling in the forward pass\n        kernel. This only applies to the `\"blackwell-fna\"` backend.\n\n    kernel_schedule (Optional[str]): Kernel type (Hopper architecture only). Choices are\n        `None`: pick the default, `\"non\"` (non-persistent), `\"coop\"` (warp-specialized\n        cooperative), or `\"pp\"` (warp-specialized ping-ponging). Refer to\n        [Hopper FMHA/FNA backend](backends.md#hopper-fna-fmha) for more information.\n\n    torch_compile (bool): Applies only to the `\"flex-fna\"` backend. Whether or not to JIT\n        compile the attention kernel. Due to this being an experimental feature in PyTorch, we\n        do not recommend it, and it is guarded by context flags. Read more in\n        [Flex Attention + `torch.compile`](context.md#flex-attention-torchcompile).\n\n    attention_kwargs: arguments to the [attention][natten.attention] operator, if used to\n        implement neighborhood cross-attention, or self attention as a fast path for\n        neighborhood attention.\n\n        If `additional_{keys,values}` are specified, NATTEN usually performs a separate\n        cross-attention using our [attention][natten.attention] operator, and\n        [merges][natten.merge_attentions] the results.\n\n        If for a given use case, the neighborhood attention problem is equivalent to self\n        attention (not causal, `kernel_size == seqlen`), NATTEN will also attempt to directly\n        use [attention][natten.attention].\n\n        You can override arguments to [attention][natten.attention] by passing a\n        dictionary here.\n\n        !!! example\n            ```python\n            out = na1d(\n                q, k, v, kernel_size=kernel_size,\n                ...,\n                attention_kwargs={\n                    \"backend\": \"blackwell-fmha\",\n                    \"run_persistent_kernel\": True,\n                }\n            )\n            ```\n\n    return_lse (bool): Whether or not to return the `logsumexp` tensor. `logsumexp` can be used\n        in the backward pass, and for [attention merging][natten.merge_attentions].\n\nReturns:\n    output (Tensor): 4-D output tensor, with the heads last layout\n        (`[batch, seqlen, heads, head_dim_v]`).\n\n    logsumexp (Tensor): only returned when `return_lse=True`. 3-D logsumexp tensor, with the\n        heads last layout (`[batch, seqlen, heads]`).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "query",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "key",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "value",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "kernel_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "int | Tuple[int]"
+          },
+          {
+            "name": "stride",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "1",
+            "annotation": "int | Tuple[int]"
+          },
+          {
+            "name": "dilation",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "1",
+            "annotation": "int | Tuple[int]"
+          },
+          {
+            "name": "is_causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool | Tuple[bool] | None"
+          },
+          {
+            "name": "scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "float | None"
+          },
+          {
+            "name": "additional_keys",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "additional_values",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "attention_kwargs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "Dict | None"
+          },
+          {
+            "name": "backend",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "str | None"
+          },
+          {
+            "name": "q_tile_shape",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "Tuple[int] | None"
+          },
+          {
+            "name": "kv_tile_shape",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "Tuple[int] | None"
+          },
+          {
+            "name": "backward_q_tile_shape",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "Tuple[int] | None"
+          },
+          {
+            "name": "backward_kv_tile_shape",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "Tuple[int] | None"
+          },
+          {
+            "name": "backward_kv_splits",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "Tuple[int] | None"
+          },
+          {
+            "name": "backward_use_pt_reduction",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "run_persistent_kernel",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "True",
+            "annotation": "bool"
+          },
+          {
+            "name": "kernel_schedule",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "str | _natten_cuda_5fe9dbb._types.KernelSchedule | None"
+          },
+          {
+            "name": "torch_compile",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "return_lse",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          }
+        ],
+        "return_annotation": "torch.Tensor | Tuple[torch.Tensor, torch.Tensor]"
+      }
+    },
+    {
+      "name": "na2d",
+      "qualname": "na2d",
+      "module": "_natten_cuda_5fe9dbb.functional",
+      "kind": "function",
+      "doc": "Computes 2-D neighborhood attention.\n\nGQA/MQA support (`heads != heads_kv`) is available. For now, `blackwell-fna` and\n`flex-fna` support GQA/MQA natively, and `cutlass-fna` and `hopper-fna` support it with\nexplicit repeats (increases memory usage and runtime).\n\nParameters:\n    query (Tensor): 2-D query tensor, with the heads last layout:\n        `[batch, X, Y, heads, head_dim]`, where token layout shape (feature map shape) is\n        `(X, Y)`.\n\n    key (Tensor): 2-D key tensor, with the heads last layout:\n        `[batch, X, Y, heads_kv, head_dim]`, where token layout shape (feature map shape) is\n        `(X, Y)`.\n\n    value (Tensor): 2-D value tensor, with the heads last layout:\n        `[batch, X, Y, heads_kv, head_dim_v]`, where token layout shape (feature map shape) is\n        `(X, Y)`.\n\n    kernel_size (Tuple[int, int] | int): Neighborhood window (kernel) size/shape. If an\n        integer, it will be repeated for all 2 dimensions. For example `kernel_size=3` is\n        reinterpreted as `kernel_size=(3, 3)`.\n\n        !!! note\n            `kernel_size` must be smaller than or equal to token layout shape (`(X, Y)`) along\n            every dimension.\n\n    stride (Tuple[int, int] | int): Sliding window step size/shape. Defaults to `1` (standard\n        sliding window). If an integer, it will be repeated for all 2 dimensions. For example\n        `stride=2` is reinterpreted as `stride=(2, 2)`.\n\n        !!! note\n            `stride` must be smaller than or equal to `kernel_size` along every dimension.\n            When `stride == kernel_size`, there will be no overlap between sliding windows,\n            which is equivalent to blocked attention (a.k.a.\n            [window self attention](https://arxiv.org/abs/2103.14030)).\n\n    dilation (Tuple[int, int] | int): Dilation step size/shape. Defaults to `1` (standard\n        sliding window). If an integer, it will be repeated for all 2 dimensions. For example\n        `dilation=4` is reinterpreted as `dilation=(4, 4)`.\n\n        !!! note\n            The product of `dilation` and `kernel_size` must be smaller than or equal to\n            token layout shape (`(X, Y)`) along every dimension.\n\n    is_causal (Tuple[bool, bool] | bool): Toggle causal masking. Defaults to `False`\n        (bi-directional). If a boolean, it will be repeated for all 2 dimensions. For example\n        `is_causal=True` is reinterpreted as `is_causal=(True, True)`.\n\n    scale (float): Attention scale. Defaults to `head_dim ** -0.5`.\n\n    additional_keys: `None` or 4-D key tensor, with the heads last layout\n        (`[batch, seqlen_kv, heads_kv, head_dim]`), corresponding to key tokens from some\n        additional context. Used when performing neighborhood cross-attention, where `query`\n        tokens attend to their neighborhood, as well as some fixed additional set of tokens.\n\n    additional_values: `None` or 4-D value tensor, with the heads last layout\n        (`[batch, seqlen_kv, heads_kv, head_dim_v]`), corresponding to value tokens from some\n        additional context. Used when performing neighborhood cross-attention, where `query`\n        tokens attend to their neighborhood, as well as some fixed additional set of tokens.\n\n        !!! note\n            `additional_keys` and `additional_values` must both either be `Tensor`s, or both\n            `None`s, and must match in shape.\n\nOther Parameters:\n    backend (str): Backend implementation to run with. Choices are: `None` (pick the best\n        available one), `\"cutlass-fna\"`, `\"hopper-fna\"`, `\"blackwell-fna\"`, `\"flex-fna\"`.\n        Refer to [backends](backends.md) for more information.\n\n    q_tile_shape (Tuple[int, int]): 2-D Tile shape for the query token layout in the forward\n        pass kernel. You can use [profiler](profiler.md) to find valid choices for your use\n        case, and search for the best combination.\n\n    kv_tile_shape (Tuple[int, int]): 2-D Tile shape for the key-value token layout in the\n        forward pass kernel. You can use [profiler](profiler.md) to find valid choices for your\n        use case, and search for the best combination.\n\n    backward_q_tile_shape (Tuple[int, int]): 2-D Tile shape for the query token layout in the\n        backward pass kernel. This is ignored by `\"flex-fna\"`.\n        You can use [profiler](profiler.md) to find valid choices for your use case, and\n        search for the best combination.\n\n    backward_kv_tile_shape (Tuple[int, int]): 2-D Tile shape for the key/value token layout in\n        the backward pass kernel. This is ignored by `\"flex-fna\"`.\n        You can use [profiler](profiler.md) to find valid choices for your use case, and\n        search for the best combination.\n\n    backward_kv_splits (Tuple[int, int]): Number of key/value tiles allowed to work in parallel\n        in the backward pass kernel. Like tile shapes, this is a tuple and not an integer for\n        neighborhood attention operations, and the size of the tuple corresponds to the number\n        of dimensions / rank of the layout of tokens. This is only respected by the\n        `\"cutlass-fna\"` backend, and only when\n        [KV parallelism](context.md#kv-parallelism-in-fna) is enabled.\n\n    backward_use_pt_reduction (bool): Whether to use PyTorch eager for computing the `dO * O`\n        product required by the backward pass, over the CUTLASS kernel. This only applies to\n        the `\"cutlass-fna\"` backend.\n\n    run_persistent_kernel (bool): Whether to use persistent tile scheduling in the forward pass\n        kernel. This only applies to the `\"blackwell-fna\"` backend.\n\n    kernel_schedule (Optional[str]): Kernel type (Hopper architecture only). Choices are\n        `None`: pick the default, `\"non\"` (non-persistent), `\"coop\"` (warp-specialized\n        cooperative), or `\"pp\"` (warp-specialized ping-ponging). Refer to\n        [Hopper FMHA/FNA backend](backends.md#hopper-fna-fmha) for more information.\n\n    torch_compile (bool): Applies only to the `\"flex-fna\"` backend. Whether or not to JIT\n        compile the attention kernel. Due to this being an experimental feature in PyTorch, we\n        do not recommend it, and it is guarded by context flags. Read more in\n        [Flex Attention + `torch.compile`](context.md#flex-attention-torchcompile).\n\n    attention_kwargs: arguments to the [attention][natten.attention] operator, if used to\n        implement neighborhood cross-attention, or self attention as a fast path for\n        neighborhood attention.\n\n        If `additional_{keys,values}` are specified, NATTEN usually performs a separate\n        cross-attention using our [attention][natten.attention] operator, and\n        [merges][natten.merge_attentions] the results.\n\n        If for a given use case, the neighborhood attention problem is equivalent to self\n        attention (not causal along any dims, `kernel_size == (X, Y)`), NATTEN will also\n        attempt to directly use [attention][natten.attention].\n\n        You can override arguments to [attention][natten.attention] by passing a\n        dictionary here.\n\n        !!! example\n            ```python\n            out = na2d(\n                q, k, v, kernel_size=kernel_size,\n                ...,\n                attention_kwargs={\n                    \"backend\": \"blackwell-fmha\",\n                    \"run_persistent_kernel\": True,\n                }\n            )\n            ```\n\n    return_lse (bool): Whether or not to return the `logsumexp` tensor. `logsumexp` can be used\n        in the backward pass, and for [attention merging][natten.merge_attentions].\n\nReturns:\n    output (Tensor): 5-D output tensor, with the heads last layout\n        (`[batch, X, Y, heads, head_dim_v]`).\n\n    logsumexp (Tensor): only returned when `return_lse=True`. 4-D logsumexp tensor, with the\n        heads last layout (`[batch, X, Y, heads]`).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "query",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "key",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "value",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "kernel_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "int | Tuple[int, int]"
+          },
+          {
+            "name": "stride",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "1",
+            "annotation": "int | Tuple[int, int]"
+          },
+          {
+            "name": "dilation",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "1",
+            "annotation": "int | Tuple[int, int]"
+          },
+          {
+            "name": "is_causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool | Tuple[bool, bool] | None"
+          },
+          {
+            "name": "scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "float | None"
+          },
+          {
+            "name": "additional_keys",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "additional_values",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "attention_kwargs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "Dict | None"
+          },
+          {
+            "name": "backend",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "str | None"
+          },
+          {
+            "name": "q_tile_shape",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "Tuple[int, int] | None"
+          },
+          {
+            "name": "kv_tile_shape",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "Tuple[int, int] | None"
+          },
+          {
+            "name": "backward_q_tile_shape",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "Tuple[int, int] | None"
+          },
+          {
+            "name": "backward_kv_tile_shape",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "Tuple[int, int] | None"
+          },
+          {
+            "name": "backward_kv_splits",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "Tuple[int, int] | None"
+          },
+          {
+            "name": "backward_use_pt_reduction",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "run_persistent_kernel",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "True",
+            "annotation": "bool"
+          },
+          {
+            "name": "kernel_schedule",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "str | _natten_cuda_5fe9dbb._types.KernelSchedule | None"
+          },
+          {
+            "name": "torch_compile",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "return_lse",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          }
+        ],
+        "return_annotation": "torch.Tensor | Tuple[torch.Tensor, torch.Tensor]"
+      }
+    },
+    {
+      "name": "na3d",
+      "qualname": "na3d",
+      "module": "_natten_cuda_5fe9dbb.functional",
+      "kind": "function",
+      "doc": "Computes 3-D neighborhood attention.\n\nGQA/MQA support (`heads != heads_kv`) is available. For now, `blackwell-fna` and\n`flex-fna` support GQA/MQA natively, and `cutlass-fna` and `hopper-fna` support it with\nexplicit repeats (increases memory usage and runtime).\n\nParameters:\n    query (Tensor): 3-D query tensor, with the heads last layout:\n        `[batch, X, Y, Z, heads, head_dim]`, where token layout shape (feature map shape) is\n        `(X, Y, Z)`.\n\n    key (Tensor): 3-D key tensor, with the heads last layout:\n        `[batch, X, Y, Z, heads_kv, head_dim]`, where token layout shape (feature map shape) is\n        `(X, Y, Z)`.\n\n    value (Tensor): 3-D value tensor, with the heads last layout:\n        `[batch, X, Y, Z, heads_kv, head_dim_V]`, where token layout shape (feature map shape) is\n        `(X, Y, Z)`.\n\n    kernel_size (Tuple[int, int, int] | int): Neighborhood window (kernel) size/shape. If an\n        integer, it will be repeated for all 3 dimensions. For example `kernel_size=3` is\n        reinterpreted as `kernel_size=(3, 3, 3)`.\n\n        !!! note\n            `kernel_size` must be smaller than or equal to token layout shape (`(X, Y, Z)`)\n            along every dimension.\n\n    stride (Tuple[int, int, int] | int): Sliding window step size/shape. Defaults to `1`\n        (standard sliding window). If an integer, it will be repeated for all 3 dimensions.\n        For example `stride=2` is reinterpreted as `stride=(2, 2, 2)`.\n\n        !!! note\n            `stride` must be smaller than or equal to `kernel_size` along every dimension.\n            When `stride == kernel_size`, there will be no overlap between sliding windows,\n            which is equivalent to blocked attention (a.k.a.\n            [window self attention](https://arxiv.org/abs/2103.14030)).\n\n    dilation (Tuple[int, int, int] | int): Dilation step size/shape. Defaults to `1` (standard\n        sliding window). If an integer, it will be repeated for all 3 dimensions. For example\n        `dilation=4` is reinterpreted as `dilation=(4, 4, 4)`.\n\n        !!! note\n            The product of `dilation` and `kernel_size` must be smaller than or equal to\n            token layout shape (`(X, Y, Z)`) along every dimension.\n\n    is_causal (Tuple[bool, bool, bool] | bool): Toggle causal masking. Defaults to `False`\n        (bi-directional). If a boolean, it will be repeated for all 3 dimensions. For example\n        `is_causal=True` is reinterpreted as `is_causal=(True, True, True)`.\n\n    scale (float): Attention scale. Defaults to `head_dim ** -0.5`.\n\n    additional_keys: `None` or 4-D key tensor, with the heads last layout\n        (`[batch, seqlen_kv, heads_kv, head_dim_v]`), corresponding to key tokens from some\n        additional context. Used when performing neighborhood cross-attention, where `query`\n        tokens attend to their neighborhood, as well as some fixed additional set of tokens.\n\n    additional_values: `None` or 4-D value tensor, with the heads last layout\n        (`[batch, seqlen_kv, heads_kv, head_dim_v]`), corresponding to value tokens from some\n        additional context. Used when performing neighborhood cross-attention, where `query`\n        tokens attend to their neighborhood, as well as some fixed additional set of tokens.\n\n        !!! note\n            `additional_keys` and `additional_values` must both either be `Tensor`s, or both\n            `None`s, and must match in shape.\n\nOther Parameters:\n    backend (str): Backend implementation to run with. Choices are: `None` (pick the best\n        available one), `\"cutlass-fna\"`, `\"hopper-fna\"`, `\"blackwell-fna\"`, `\"flex-fna\"`.\n        Refer to [backends](backends.md) for more information.\n\n    q_tile_shape (Tuple[int, int, int]): 3-D Tile shape for the query token layout in the\n        forward pass kernel. You can use [profiler](profiler.md) to find valid choices for your\n        use case, and search for the best combination.\n\n    kv_tile_shape (Tuple[int, int, int]): 3-D Tile shape for the key-value token layout in the\n        forward pass kernel. You can use [profiler](profiler.md) to find valid choices for your\n        use case, and search for the best combination.\n\n    backward_q_tile_shape (Tuple[int, int, int]): 3-D Tile shape for the query token layout in\n        the backward pass kernel. This is ignored by `\"flex-fna\"`.\n        You can use [profiler](profiler.md) to find valid choices for your use case, and\n        search for the best combination.\n\n    backward_kv_tile_shape (Tuple[int, int, int]): 3-D Tile shape for the key/value token\n        layout in the backward pass kernel. This is ignored by `\"flex-fna\"`.\n        You can use [profiler](profiler.md) to find valid choices for your use case,\n        and search for the best combination.\n\n    backward_kv_splits (Tuple[int, int, int]): Number of key/value tiles allowed to work in\n        parallel in the backward pass kernel. Like tile shapes, this is a tuple and not an\n        integer for neighborhood attention operations, and the size of the tuple corresponds to\n        the number of dimensions / rank of the layout of tokens. This is only respected by the\n        `\"cutlass-fna\"` backend, and only when\n        [KV parallelism](context.md#kv-parallelism-in-fna) is enabled.\n\n    backward_use_pt_reduction (bool): Whether to use PyTorch eager for computing the `dO * O`\n        product required by the backward pass, over the CUTLASS kernel. This only applies to\n        the `\"cutlass-fna\"` backend.\n\n    run_persistent_kernel (bool): Whether to use persistent tile scheduling in the forward pass\n        kernel. This only applies to the `\"blackwell-fna\"` backend.\n\n    kernel_schedule (Optional[str]): Kernel type (Hopper architecture only). Choices are\n        `None`: pick the default, `\"non\"` (non-persistent), `\"coop\"` (warp-specialized\n        cooperative), or `\"pp\"` (warp-specialized ping-ponging). Refer to\n        [Hopper FMHA/FNA backend](backends.md#hopper-fna-fmha) for more information.\n\n    torch_compile (bool): Applies only to the `\"flex-fna\"` backend. Whether or not to JIT\n        compile the attention kernel. Due to this being an experimental feature in PyTorch, we\n        do not recommend it, and it is guarded by context flags. Read more in\n        [Flex Attention + `torch.compile`](context.md#flex-attention-torchcompile).\n\n    attention_kwargs: arguments to the [attention][natten.attention] operator, if used to\n        implement neighborhood cross-attention, or self attention as a fast path for\n        neighborhood attention.\n\n        If `additional_{keys,values}` are specified, NATTEN usually performs a separate\n        cross-attention using our [attention][natten.attention] operator, and\n        [merges][natten.merge_attentions] the results.\n\n        If for a given use case, the neighborhood attention problem is equivalent to self\n        attention (not causal along any dims, `kernel_size == (X, Y, Z)`), NATTEN will also\n        attempt to directly use [attention][natten.attention].\n\n        You can override arguments to [attention][natten.attention] by passing a\n        dictionary here.\n\n        !!! example\n            ```python\n            out = na3d(\n                q, k, v, kernel_size=kernel_size,\n                ...,\n                attention_kwargs={\n                    \"backend\": \"blackwell-fmha\",\n                    \"run_persistent_kernel\": True,\n                }\n            )\n            ```\n\n    return_lse (bool): Whether or not to return the `logsumexp` tensor. `logsumexp` can be used\n        in the backward pass, and for [attention merging][natten.merge_attentions].\n\nReturns:\n    output (Tensor): 6-D output tensor, with the heads last layout\n        (`[batch, X, Y, Z, heads, head_dim_v]`).\n\n    logsumexp (Tensor): only returned when `return_lse=True`. 5-D logsumexp tensor, with the\n        heads last layout (`[batch, X, Y, Z, heads]`).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "query",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "key",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "value",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "kernel_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "int | Tuple[int, int, int]"
+          },
+          {
+            "name": "stride",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "1",
+            "annotation": "int | Tuple[int, int, int]"
+          },
+          {
+            "name": "dilation",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "1",
+            "annotation": "int | Tuple[int, int, int]"
+          },
+          {
+            "name": "is_causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool | Tuple[bool, bool, bool] | None"
+          },
+          {
+            "name": "scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "float | None"
+          },
+          {
+            "name": "additional_keys",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "additional_values",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "attention_kwargs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "Dict | None"
+          },
+          {
+            "name": "backend",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "str | None"
+          },
+          {
+            "name": "q_tile_shape",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "Tuple[int, int, int] | None"
+          },
+          {
+            "name": "kv_tile_shape",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "Tuple[int, int, int] | None"
+          },
+          {
+            "name": "backward_q_tile_shape",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "Tuple[int, int, int] | None"
+          },
+          {
+            "name": "backward_kv_tile_shape",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "Tuple[int, int, int] | None"
+          },
+          {
+            "name": "backward_kv_splits",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "Tuple[int, int, int] | None"
+          },
+          {
+            "name": "backward_use_pt_reduction",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "run_persistent_kernel",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "True",
+            "annotation": "bool"
+          },
+          {
+            "name": "kernel_schedule",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "str | _natten_cuda_5fe9dbb._types.KernelSchedule | None"
+          },
+          {
+            "name": "torch_compile",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "return_lse",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          }
+        ],
+        "return_annotation": "torch.Tensor | Tuple[torch.Tensor, torch.Tensor]"
+      }
+    },
+    {
+      "name": "attention",
+      "qualname": "attention",
+      "module": "_natten_cuda_5fe9dbb.functional",
+      "kind": "function",
+      "doc": "Runs standard dot product attention.\n\nThis operation is used to implement neighborhood cross attention, in which we allow every\ntoken to interact with some additional context (`additional_keys` and `additional_values`\ntensors in [na1d][natten.na1d], [na2d][natten.na2d], and [na3d][natten.na3d]).\nThis operator is also used as a fast path for cases where neighborhood attention is equivalent\nto self attention (not causal along any dims, and `kernel_size` is equal to the number of input\ntokens).\n\nThis operation does not call into PyTorch's SDPA, and only runs one of the NATTEN backends\n(`cutlass-fmha`, `hopper-fmha`, `blackwell-fmha`, `flex-fmha`). Reasons for that include being\nable to control performance-related arguments, return logsumexp, and more.\nFor more information refer to [backends](backends.md).\n\nCausal mask, and Variable length (varlen) Attention are also supported in some backends\n(`cutlass-fmha` and `blackwell-fmha`).\n\nVarlen Attention is only supported for the sequence-packed layout: QKV tensors have batch size\n1, and tokens from different batches are concatenated without any padding along the sequence\ndimension. Sequence lengths for different batches can be provided in two ways:\n    1. `seqlens_Q` and `seqlens_KV` (less efficient): only provide the sequence lengths as\n        integer tensors (must be on the same device as QKV), and NATTEN will compute cumulative\n        and maximum sequence lengths on each call.\n        This is **incompatible** with full-graph `torch.compile` since it requires a\n        synchronization.\n    2. `cumulative_seqlen_{Q,KV}` and `max_seqlen_{Q,KV}` (more efficient):\n        compute cumulative and maximum sequence lengths. `cumulative_seqlen_{Q,KV}` are integer\n        tensors on the same device as QKV containing the cumulative sum of `seqlens_{Q,KV}`,\n        with an additional `0` element in the beginning, therefore sized `batch+1`.\n        `max_seqlen_{Q,KV}` are integers (not Tensors) that represent the maximum sequence\n        lengths for Q and KV among all sequence batches.\n        You can use `natten.utils.varlen.generate_varlen_parameters` to generate these\n        parameters:\n            ```python3\n            from .utils.varlen import generate_varlen_parameters\n            (\n                cumulative_seqlen_Q,\n                cumulative_seqlen_KV,\n                max_seqlen_Q,\n                max_seqlen_KV,\n            ) = generate_varlen_parameters(q, k, v, seqlens_Q, seqlens_KV)\n            ```\n        As long as `generate_varlen_parameters` is called ahead of torch.compiling the model, it\n        is supported without any graph breaks.\n\nGQA/MQA support (`heads != heads_kv`) is available. For now, `blackwell-fmha` and\n`flex-fmha` support GQA/MQA natively, and `cutlass-fmha` and `hopper-fmha` support it with\nexplicit repeats (increases memory usage and runtime).\n\nParameters:\n    query (Tensor): 4-D query tensor, with the heads last layout\n        (`[batch, seqlen, heads, head_dim]`)\n\n    key (Tensor): 4-D key tensor, with the heads last layout\n        (`[batch, seqlen_kv, heads_kv, head_dim]`)\n\n    value (Tensor): 4-D value tensor, with the heads last layout\n        (`[batch, seqlen_kv, heads_kv, head_dim_v]`)\n\n    is_causal (bool): Toggle causal masking. Defaults to `False` (bi-directional).\n\n    scale (float): Attention scale. Defaults to `head_dim ** -0.5`.\n\n    seqlens_Q (Optional[Tensor]): (varlen) Optional 1-D tensor with size `batch`\n        indicating the number of query tokens in each batch. Must be passed together with\n        `seqlens_KV`.\n\n    seqlens_KV (Optional[Tensor]): (varlen) Optional 1-D tensor with size `batch`\n        indicating the number of key/value tokens in each batch. Must be passed together with\n        `seqlens_Q`.\n\n    cumulative_seqlen_Q (Optional[Tensor]): (varlen) Optional 1-D tensor with size `batch + 1`\n        indicating the cumulative sum of number of query tokens in each batch, with an\n        additional 0 element in the beginning. Must be passed together with\n        `cumulative_seqlen_KV` and `max_seqlen_{Q,KV}`.\n\n    cumulative_seqlen_KV (Optional[Tensor]): (varlen) Optional 1-D tensor with size `batch + 1`\n        indicating the cumulative sum of number of key/value tokens in each batch, with an\n        additional 0 element in the beginning. Must be passed together with\n        `cumulative_seqlen_Q` and `max_seqlen_{Q,KV}`.\n\n    max_seqlen_Q (Optional[int]): (varlen) Optional integer indicating the maximum query\n        sequence length in all batches. Must be passed together with `cumulative_seqlen_{Q,KV}`\n        and `max_seqlen_KV`.\n\n    max_seqlen_KV (Optional[int]): (varlen) Optional integer indicating the maximum key/value\n        sequence length in all batches. Must be passed together with `cumulative_seqlen_{Q,KV}`\n        and `max_seqlen_Q`.\n\nOther Parameters:\n    backend (str): Backend implementation to run with. Choices are: `None` (pick the best\n        available one), `\"cutlass-fmha\"`, `\"hopper-fmha\"`, `\"blackwell-fmha\"`, `\"flex-fmha\"`.\n        Refer to [backends](backends.md) for more information.\n\n    q_tile_size (int): Tile size along query sequence length in the forward pass kernel.\n        You can use [profiler](profiler.md) to find valid choices for your use case.\n\n    kv_tile_size (int): Tile size along key/value sequence length in the forward pass kernel.\n        You can use [profiler](profiler.md) to find valid choices for your use case.\n\n    backward_q_tile_size (int): Tile size along query sequence length in the backward pass\n        kernel. This is ignored by `\"flex-fmha\"`.\n        You can use [profiler](profiler.md) to find valid choices for your use case.\n\n    backward_kv_tile_size (int): Tile size along key/value sequence length in the backward pass\n        kernel. This is ignored by `\"flex-fmha\"`.\n        You can use [profiler](profiler.md) to find valid choices for your use case.\n\n    backward_kv_splits (int): Number of key/value tiles allowed to work in parallel in the\n        backward pass kernel. This is only respected by the `\"cutlass-fmha\"` backend, only when\n        [KV parallelism](context.md#kv-parallelism-in-fna) is enabled.\n\n    backward_use_pt_reduction (bool): Whether to use PyTorch eager for computing the `dO * O`\n        product required by the backward pass, over the CUTLASS kernel. This only applies to\n        the `\"cutlass-fmha\"` backend.\n\n    run_persistent_kernel (bool): Whether to use persistent tile scheduling in the forward pass\n        kernel. This only applies to the `\"blackwell-fmha\"` backend.\n\n    kernel_schedule (Optional[str]): Kernel type (Hopper architecture only). Choices are\n        `None`: pick the default, `\"non\"` (non-persistent), `\"coop\"` (warp-specialized\n        cooperative), or `\"pp\"` (warp-specialized ping-ponging). Refer to\n        [Hopper FMHA/FNA backend](backends.md#hopper-fna-fmha) for more information.\n\n    torch_compile (bool): Applies only to the `\"flex-fmha\"` backend. Whether or not to JIT\n        compile the attention kernel. Due to this being an experimental feature in PyTorch, we\n        do not recommend it, and it is guarded by context flags. Read more in\n        [Flex Attention + `torch.compile`](context.md#flex-attention-torchcompile).\n\n    return_lse (bool): Whether or not to return the `logsumexp` tensor. `logsumexp` can be used\n        in the backward pass, and for [attention merging][natten.merge_attentions].\n\nReturns:\n    output (Tensor): 4-D output tensor, with the heads last layout\n        (`[batch, seqlen, heads, head_dim_v]`).\n\n    logsumexp (Tensor): only returned when `return_lse=True`. 3-D logsumexp tensor, with the\n        heads last layout (`[batch, seqlen, heads]`).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "query",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "key",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "value",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "is_causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "float | None"
+          },
+          {
+            "name": "seqlens_Q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "seqlens_KV",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "cumulative_seqlen_Q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "cumulative_seqlen_KV",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "max_seqlen_Q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "int | None"
+          },
+          {
+            "name": "max_seqlen_KV",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "int | None"
+          },
+          {
+            "name": "backend",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "str | None"
+          },
+          {
+            "name": "q_tile_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "int | None"
+          },
+          {
+            "name": "kv_tile_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "int | None"
+          },
+          {
+            "name": "backward_q_tile_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "int | None"
+          },
+          {
+            "name": "backward_kv_tile_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "int | None"
+          },
+          {
+            "name": "backward_kv_splits",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "int | None"
+          },
+          {
+            "name": "backward_use_pt_reduction",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "run_persistent_kernel",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "True",
+            "annotation": "bool"
+          },
+          {
+            "name": "kernel_schedule",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "str | _natten_cuda_5fe9dbb._types.KernelSchedule | None"
+          },
+          {
+            "name": "torch_compile",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "return_lse",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          }
+        ],
+        "return_annotation": "torch.Tensor | Tuple[torch.Tensor, torch.Tensor]"
+      }
+    },
+    {
+      "name": "merge_attentions",
+      "qualname": "merge_attentions",
+      "module": "_natten_cuda_5fe9dbb.attn_merge",
+      "kind": "function",
+      "doc": "Takes multiple attention *outputs* originating from the same query tensor, and their\ncorresponding logsumexps, and merges them as if their context (key/value pair) had been\nconcatenated.\n\nThis operation is used to implement cross-neighborhood attention, and can also be used for\ndistributed setups, such as context-parallelism.\n\nThis operation also attempts to use `torch.compile` to fuse the elementwise operations. This\ncan be disabled by passing `torch_compile=False`.\n\nParameters:\n    outputs (List[Tensor]): List of 4-D attention output tensors, with the heads last layout\n        (`[batch, seqlen, heads, head_dim]`)\n\n    lse_tensors (List[Tensor]): List of 3-D logsumexp tensors, with the heads last layout\n        (`[batch, seqlen, heads]`)\n\n    torch_compile (bool): Attempt to use `torch.compile` to fuse the underlying elementwise\n        operations. Default: True.\n\n    use_autograd_fix (bool): fix backpropagation by using a custom autograd function. Only\n        compatible with fused attention operations (Flash/FMHA/FNA), only as long as the inputs\n        of this function are (views) of outputs from said attention operation.\n        NATTEN's tests (tests/test_attn_merge.py) only verify correctness for when using\n        attention operations from NATTEN. Integration for non-NATTEN ops must be verified by the\n        end user.\n        This must be disabled when using unfused Attention, which includes Flex without\n        torch.compile. Default: True.\n\nReturns:\n    output (Tensor): merged attention output.\n\n    logsumexp (Tensor): updated logsumexp.",
+      "signature": {
+        "parameters": [
+          {
+            "name": "outputs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "List[torch.Tensor]"
+          },
+          {
+            "name": "lse_tensors",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "List[torch.Tensor]"
+          },
+          {
+            "name": "torch_compile",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "True",
+            "annotation": "bool"
+          },
+          {
+            "name": "use_autograd_fix",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "True",
+            "annotation": "bool"
+          }
+        ],
+        "return_annotation": "Tuple[torch.Tensor, torch.Tensor]"
+      }
+    }
+  ],
+  "layers": []
+}

build/torch210-cu126-aarch64-linux/token_permute/__init__.py ADDED Viewed

	@@ -0,0 +1,32 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+from ..token_permute.frontend import (
+    token_permute_operation,
+    token_unpermute_operation,
+)
+__all__ = [
+    "token_permute_operation",
+    "token_unpermute_operation",
+]

build/torch210-cu126-aarch64-linux/token_permute/cutlass_impl.py ADDED Viewed

	@@ -0,0 +1,286 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+import functools
+from typing import Tuple
+import torch
+from torch import Tensor
+from torch.amp import custom_bwd, custom_fwd
+from torch.autograd import Function
+amp_fwd = functools.partial(custom_fwd, device_type="cuda")
+amp_bwd = functools.partial(custom_bwd, device_type="cuda")
+from .._libnatten import (
+    HAS_LIBNATTEN,
+    token_permute_1d,
+    token_permute_2d,
+    token_permute_3d,
+    token_unpermute_1d,
+    token_unpermute_2d,
+    token_unpermute_3d,
+)
+from .._types import DimensionType, NoneType
+from ..utils import log
+from ..utils.device import get_device_cc, is_cuda
+logger = log.get_logger(__name__)
+def can_run_cutlass_tokperm(tensor: Tensor) -> bool:
+    if not HAS_LIBNATTEN:
+        logger.debug(
+            "Can't use libnatten TokPerm kernels, because libnatten is not available."
+        )
+        return False
+    if not is_cuda(tensor.device):
+        logger.debug(
+            "Can't use libnatten TokPerm kernels, because input is not a CUDA tensor."
+        )
+        return False
+    is_fp8_allowed = get_device_cc(tensor.device) in [100, 103]
+    if tensor.dtype not in [
+        torch.float32,
+        torch.float16,
+        torch.bfloat16,
+        torch.float16,
+    ] and (
+        is_fp8_allowed and tensor.dtype not in [torch.float8_e5m2, torch.float8_e4m3fn]
+    ):
+        logger.debug(
+            f"Can't use libnatten TokPerm kernels; unexpected dtype {tensor.dtype}."
+        )
+        return False
+    return True
+PERMUTE_OPS = {1: token_permute_1d, 2: token_permute_2d, 3: token_permute_3d}
+UNPERMUTE_OPS = {1: token_unpermute_1d, 2: token_unpermute_2d, 3: token_unpermute_3d}
+def make_cutlass_token_permute_autograd_fn(na_dim):
+    assert na_dim in [1, 2, 3]
+    class CutlassTokenPermuteAutogradFn(Function):
+        @staticmethod
+        @amp_fwd
+        def forward(
+            ctx,
+            tensor: Tensor,
+            tile_shape: DimensionType,
+            dilation: DimensionType,
+            flip_tiled_dims: bool,
+        ) -> Tensor:
+            output = PERMUTE_OPS[na_dim](
+                tensor,
+                tile_shape=tile_shape,
+                dilation=dilation,
+                flip_tiled_dims=flip_tiled_dims,
+            )
+            ctx.tile_shape = tile_shape
+            ctx.dilation = dilation
+            ctx.flip_tiled_dims = flip_tiled_dims
+            ctx.token_layout = tuple(x for x in tensor.shape[1 : na_dim + 1])
+            assert len(ctx.token_layout) == na_dim
+            return output
+        @staticmethod
+        @amp_bwd
+        def backward(ctx, d_output: Tensor) -> Tuple[
+            Tensor,
+            NoneType,
+            NoneType,
+            NoneType,
+        ]:
+            d_output_unpermuted = UNPERMUTE_OPS[na_dim](
+                d_output,
+                token_layout_shape=ctx.token_layout,
+                tile_shape=ctx.tile_shape,
+                dilation=ctx.dilation,
+                flip_tiled_dims=ctx.flip_tiled_dims,
+            )
+            return (
+                d_output_unpermuted,
+                None,
+                None,
+                None,
+            )
+    return CutlassTokenPermuteAutogradFn
+def make_cutlass_token_unpermute_autograd_fn(na_dim):
+    assert na_dim in [1, 2, 3]
+    class CutlassTokenUnPermuteAutogradFn(Function):
+        @staticmethod
+        @amp_fwd
+        def forward(
+            ctx,
+            tensor: Tensor,
+            token_layout: DimensionType,
+            tile_shape: DimensionType,
+            dilation: DimensionType,
+            flip_tiled_dims: bool,
+        ) -> Tensor:
+            output = UNPERMUTE_OPS[na_dim](
+                tensor,
+                token_layout_shape=token_layout,
+                tile_shape=tile_shape,
+                dilation=dilation,
+                flip_tiled_dims=flip_tiled_dims,
+            )
+            ctx.tile_shape = tile_shape
+            ctx.dilation = dilation
+            ctx.flip_tiled_dims = flip_tiled_dims
+            return output
+        @staticmethod
+        @amp_bwd
+        def backward(ctx, d_output: Tensor) -> Tuple[
+            Tensor,
+            NoneType,
+            NoneType,
+            NoneType,
+            NoneType,
+        ]:
+            d_output_permuted = PERMUTE_OPS[na_dim](
+                d_output,
+                tile_shape=ctx.tile_shape,
+                dilation=ctx.dilation,
+                flip_tiled_dims=ctx.flip_tiled_dims,
+            )
+            return (
+                d_output_permuted,
+                None,
+                None,
+                None,
+                None,
+            )
+    return CutlassTokenUnPermuteAutogradFn
+CutlassTokenPermute1DAutogradFn = make_cutlass_token_permute_autograd_fn(1)
+CutlassTokenPermute2DAutogradFn = make_cutlass_token_permute_autograd_fn(2)
+CutlassTokenPermute3DAutogradFn = make_cutlass_token_permute_autograd_fn(3)
+CutlassTokenUnPermute1DAutogradFn = make_cutlass_token_unpermute_autograd_fn(1)
+CutlassTokenUnPermute2DAutogradFn = make_cutlass_token_unpermute_autograd_fn(2)
+CutlassTokenUnPermute3DAutogradFn = make_cutlass_token_unpermute_autograd_fn(3)
+CutlassTokenPermuteAutogradFns = {
+    1: CutlassTokenPermute1DAutogradFn,
+    2: CutlassTokenPermute2DAutogradFn,
+    3: CutlassTokenPermute3DAutogradFn,
+}
+CutlassTokenUnPermuteAutogradFns = {
+    1: CutlassTokenUnPermute1DAutogradFn,
+    2: CutlassTokenUnPermute2DAutogradFn,
+    3: CutlassTokenUnPermute3DAutogradFn,
+}
+def token_permute_cutlass(
+    tensor: Tensor,
+    tile_shape: DimensionType,
+    dilation: DimensionType,
+    flip_tiled_dims,
+) -> Tensor:
+    if tensor.dim() not in [4, 5, 6]:
+        raise ValueError(
+            "Expected 4D, 5D, or 6D tensor (corresponding to NA1D, 2D, 3D), "
+            f"got {tensor.dim()}D input."
+        )
+    na_dim = tensor.dim() - 3
+    assert na_dim in [1, 2, 3]
+    if len(tile_shape) != na_dim:
+        raise ValueError(
+            f"Expected {na_dim}D tiler for NA{na_dim}D, " f"got {tile_shape=}."
+        )
+    if not can_run_cutlass_tokperm(tensor):
+        raise NotImplementedError(
+            "Use case is not compatible with CUTLASS Token Permute."
+        )
+    tensor = tensor.contiguous()
+    output = CutlassTokenPermuteAutogradFns[na_dim].apply(
+        tensor,
+        tile_shape,
+        dilation,
+        flip_tiled_dims,
+    )
+    return output
+def token_unpermute_cutlass(
+    tensor: Tensor,
+    token_layout_shape: DimensionType,
+    tile_shape: DimensionType,
+    dilation: DimensionType,
+    flip_tiled_dims: bool,
+) -> Tensor:
+    if tensor.dim() != 4:
+        raise ValueError(f"Expected flattened 4D tensor, got {tensor.dim()}D input.")
+    na_dim = len(token_layout_shape)
+    if len(tile_shape) != na_dim:
+        raise ValueError(
+            f"Expected {na_dim}D tiler for NA{na_dim}D, " f"got {tile_shape=}."
+        )
+    if not can_run_cutlass_tokperm(tensor):
+        raise NotImplementedError(
+            "Use case is not compatible with CUTLASS Token UnPermute."
+        )
+    tensor = tensor.contiguous()
+    output = CutlassTokenUnPermuteAutogradFns[na_dim].apply(
+        tensor,
+        token_layout_shape,
+        tile_shape,
+        dilation,
+        flip_tiled_dims,
+    )
+    return output

build/torch210-cu126-aarch64-linux/token_permute/frontend.py ADDED Viewed

	@@ -0,0 +1,137 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+from typing import Optional
+from torch import Tensor
+from .._environment import USE_TORCH_IMPL_DEFAULT
+from ..token_permute.cutlass_impl import (
+    can_run_cutlass_tokperm,
+    token_permute_cutlass,
+    token_unpermute_cutlass,
+)
+from ..token_permute.torch_impl import token_permute_torch, token_unpermute_torch
+from .._types import DimensionType
+from ..utils import log
+from ..utils.tuples import ceil_div_tuple, mul_tuple
+logger = log.get_logger(__name__)
+def token_permute_operation(
+    tensor: Tensor,
+    tile_shape: DimensionType,
+    dilation: Optional[DimensionType] = None,
+    flip_tiled_dims: bool = True,
+    use_torch: bool = USE_TORCH_IMPL_DEFAULT,
+) -> tuple[Tensor, DimensionType, DimensionType]:
+    if tensor.dim() not in [4, 5, 6]:
+        raise ValueError(
+            "Expected 4D, 5D, or 6D tensor (corresponding to NA1D, 2D, 3D), "
+            f"got {tensor.dim()}D input."
+        )
+    na_dim = tensor.dim() - 3
+    assert na_dim in [1, 2, 3]
+    if len(tile_shape) != na_dim:
+        raise ValueError(
+            f"Expected {na_dim}D tiler for NA{na_dim}D, " f"got {tile_shape=}."
+        )
+    if dilation is not None and len(dilation) != na_dim:
+        raise ValueError(
+            f"Expected {na_dim}D dilation for NA{na_dim}D, " f"got {dilation=}."
+        )
+    dilation_: DimensionType = dilation or tuple(1 for _ in range(na_dim))  # type: ignore[assignment]
+    tensor = tensor.contiguous()
+    batch, *token_layout_, heads, dim = tensor.shape
+    token_layout: DimensionType = tuple(x for x in token_layout_)  # type: ignore[assignment]
+    token_layout_post_dilation: DimensionType = mul_tuple(ceil_div_tuple(ceil_div_tuple(token_layout, tile_shape), dilation_), tile_shape)  # type: ignore[assignment]
+    if not use_torch and can_run_cutlass_tokperm(tensor):
+        output = token_permute_cutlass(
+            tensor,
+            tile_shape=tile_shape,
+            dilation=dilation_,
+            flip_tiled_dims=flip_tiled_dims,
+        )
+    else:
+        output = token_permute_torch(
+            tensor,
+            tile_shape=tile_shape,
+            dilation=dilation_,
+            flip_tiled_dims=flip_tiled_dims,
+        )
+    return output, token_layout, token_layout_post_dilation
+def token_unpermute_operation(
+    tensor: Tensor,
+    token_layout_shape: DimensionType,
+    tile_shape: DimensionType,
+    dilation: Optional[DimensionType] = None,
+    flip_tiled_dims: bool = True,
+    use_torch: bool = USE_TORCH_IMPL_DEFAULT,
+) -> Tensor:
+    if tensor.dim() != 4:
+        raise ValueError(f"Expected flattened 4D tensor, got {tensor.dim()}D input.")
+    na_dim = len(token_layout_shape)
+    if len(tile_shape) != na_dim:
+        raise ValueError(
+            f"Expected {na_dim}D tiler for NA{na_dim}D, " f"got {tile_shape=}."
+        )
+    if dilation is not None and len(dilation) != na_dim:
+        raise ValueError(
+            f"Expected {na_dim}D dilation for NA{na_dim}D, " f"got {dilation=}."
+        )
+    dilation_: DimensionType = dilation or tuple(1 for _ in range(na_dim))  # type: ignore[assignment]
+    tensor = tensor.contiguous()
+    if not use_torch and can_run_cutlass_tokperm(tensor):
+        output = token_unpermute_cutlass(
+            tensor,
+            token_layout_shape,
+            tile_shape=tile_shape,
+            dilation=dilation_,
+            flip_tiled_dims=flip_tiled_dims,
+        )
+    else:
+        output = token_unpermute_torch(
+            tensor,
+            token_layout_shape,
+            tile_shape=tile_shape,
+            dilation=dilation_,
+            flip_tiled_dims=flip_tiled_dims,
+        )
+    return output

build/torch210-cu126-aarch64-linux/token_permute/torch_impl.py ADDED Viewed

	@@ -0,0 +1,368 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+import math
+import torch
+from torch import Tensor
+from .._types import DimensionType
+from ..utils import log
+from ..utils.environment import is_torch_compiling
+from ..utils.tuples import ceil_div_tuple, mul_tuple, sub_tuple
+logger = log.get_logger(__name__)
+DISABLE_PADDING_WARNING = True
+TOKEN_PERMUTE_PADDING_RATIO_LIMIT_UNTIL_WARNING = 0.5
+def _maybe_pad(
+    tensor: Tensor, tile_shape: DimensionType, dilation: DimensionType
+) -> Tensor:
+    if tensor.dim() not in [4, 5, 6]:
+        raise ValueError(
+            "Expected 4D, 5D, or 6D tensor (corresponding to NA1D, 2D, 3D), "
+            f"got {tensor.dim()}D input."
+        )
+    na_dim = tensor.dim() - 3
+    assert na_dim in [1, 2, 3]
+    if len(tile_shape) != na_dim:
+        raise ValueError(
+            f"Expected {na_dim}D tiler for NA{na_dim}D, " f"got {tile_shape=}."
+        )
+    if dilation is not None and len(dilation) != na_dim:
+        raise ValueError(
+            f"Expected {na_dim}D dilation for NA{na_dim}D, " f"got {dilation=}."
+        )
+    token_layout = tensor.shape[1 : na_dim + 1]
+    tile_shape_ = tuple(x for x in tile_shape)
+    if dilation is not None:
+        # NOTE: LCM?
+        # tile_shape_ = tuple(math.lcm(t, d) for t, d in zip(tile_shape, dilation))
+        tile_shape_ = tuple(t * d for t, d in zip(tile_shape, dilation))
+    rest = tuple((x + t - 1) // t for x, t in zip(token_layout, tile_shape_))
+    residual = tuple(r * t - x for x, t, r in zip(token_layout, tile_shape_, rest))
+    assert all(res >= 0 for res in residual)
+    if not DISABLE_PADDING_WARNING and any(
+        res / sz > TOKEN_PERMUTE_PADDING_RATIO_LIMIT_UNTIL_WARNING
+        for res, sz in zip(residual, token_layout)
+    ):
+        padded_token_layout = tuple(x + p for x, p in zip(token_layout, residual))
+        logger.warning(
+            "Potentially excessive padding detected in token permute: "
+            f"input shape {token_layout} will be padded to {padded_token_layout} to handle "
+            "token permutation, which can result in excessive memory usage, and "
+            "performance implications. Consider choosing your tile shapes, input shapes "
+            "(and dilation if you use it) accordingly. Refer to NATTEN docs for more info."
+        )
+    if any(res > 0 for res in residual):
+        padding = [0, 0, 0, 0]  # head_dim_left, head_dim_right, heads_left, heads_right
+        for res in reversed(residual):
+            padding.append(0)  # left pad
+            padding.append(res)  # right pad
+        tensor_padded = torch.nn.functional.pad(tensor, padding, "constant", 0)
+    else:
+        tensor_padded = tensor
+    return tensor_padded
+def _token_permute(
+    tensor: Tensor,
+    tile_shape: DimensionType,
+    dilation: DimensionType,
+    flip_tiled_dims: bool,
+) -> Tensor:
+    if tensor.dim() not in [4, 5, 6]:
+        raise ValueError(
+            "Expected 4D, 5D, or 6D tensor (corresponding to NA1D, 2D, 3D), "
+            f"got {tensor.dim()}D input."
+        )
+    na_dim = tensor.dim() - 3
+    assert na_dim in [1, 2, 3]
+    if len(tile_shape) != na_dim:
+        raise ValueError(
+            f"Expected {na_dim}D tiler for NA{na_dim}D, " f"got {tile_shape=}."
+        )
+    dilation = dilation or tuple(1 for _ in range(na_dim))
+    if len(dilation) != na_dim:
+        raise ValueError(
+            f"Expected {na_dim}D dilation for NA{na_dim}D, " f"got {dilation=}."
+        )
+    batch, *token_layout, heads, dim = tensor.shape
+    if any(
+        x % d != 0 or (x // d) % t != 0
+        for x, t, d in zip(token_layout, tile_shape, dilation)
+    ):
+        raise ValueError(
+            "Tensor must be divisible by static tile shape and dilation, but got "
+            f"{tensor.shape=}, {tile_shape=}, {dilation=}."
+        )
+    num_dilation_groups = math.prod(dilation)
+    token_layout_post_dilation = tuple(x // d for x, d in zip(token_layout, dilation))
+    rest = tuple(x // d // t for x, t, d in zip(token_layout, tile_shape, dilation))
+    logical_divide_dims = []
+    for d, r, t in zip(dilation, rest, tile_shape):
+        logical_divide_dims += [r, t, d]
+    # Two permutations at once:
+    #  1. logical divide to tiled divide
+    #  2. (optionally) flip order of tiled modes (i.e. (X,Y,Z) -> (Z,Y,X)) for compatibility with
+    #     CuTe's identity layout mapping.
+    permutation_idxes_r = []
+    permutation_idxes_t = []
+    permutation_idxes_d = []
+    for i in range(na_dim):
+        if flip_tiled_dims:
+            permutation_idxes_r += [(na_dim - i - 1) * 3 + 1]
+            permutation_idxes_t += [(na_dim - i - 1) * 3 + 2]
+            permutation_idxes_d += [(na_dim - i - 1) * 3 + 3]
+        else:
+            permutation_idxes_r += [i * 3 + 1]
+            permutation_idxes_t += [i * 3 + 2]
+            permutation_idxes_d += [i * 3 + 3]
+    permutation_idxes = (
+        [0]
+        + permutation_idxes_d
+        + permutation_idxes_r
+        + permutation_idxes_t
+        + [na_dim * 3 + 1, na_dim * 3 + 2]
+    )
+    # View, not copy
+    tensor_tiled = tensor.view(batch, *logical_divide_dims, heads, dim)
+    if not is_torch_compiling():
+        assert tensor_tiled.data_ptr() == tensor.data_ptr()
+    # View, not copy
+    tensor_permuted = tensor_tiled.permute(*permutation_idxes)
+    if not is_torch_compiling():
+        assert tensor_permuted.data_ptr() == tensor_tiled.data_ptr()
+    # Reshape back and copy
+    tensor_flatten = tensor_permuted.reshape(
+        num_dilation_groups * batch, math.prod(token_layout_post_dilation), heads, dim
+    ).contiguous()
+    # NOTE: token permute without dilation is a no-op for 1-D
+    # assert na_dim == 1 or tensor_flatten.data_ptr() != tensor_permuted.data_ptr()
+    assert tensor_flatten.is_contiguous()
+    return tensor_flatten
+def _token_unpermute(
+    tensor: Tensor,
+    token_layout: DimensionType,
+    tile_shape: DimensionType,
+    dilation: DimensionType,
+    flip_tiled_dims: bool,
+):
+    if tensor.dim() != 4:
+        raise ValueError(f"Expected flattened 4D tensor, got {tensor.dim()}D input.")
+    na_dim = len(token_layout)
+    assert na_dim in [1, 2, 3]
+    if len(tile_shape) != na_dim:
+        raise ValueError(
+            f"Expected {na_dim}D tiler for NA{na_dim}D, " f"got {tile_shape=}."
+        )
+    dilation = dilation or tuple(1 for _ in range(na_dim))
+    if len(dilation) != na_dim:
+        raise ValueError(
+            f"Expected {na_dim}D dilation for NA{na_dim}D, " f"got {dilation=}."
+        )
+    num_dilation_groups = math.prod(dilation)
+    batch, seqlen, heads, dim = tensor.shape
+    if batch % num_dilation_groups != 0:
+        raise ValueError(
+            "Expected batch size in token-permuted tensor to be divisible by "
+            f"number of dilation groups {num_dilation_groups} ({dilation=}), got {batch=}."
+        )
+    batch_actual = batch // num_dilation_groups
+    rest_shape = ceil_div_tuple(ceil_div_tuple(token_layout, tile_shape), dilation)
+    token_layout_padded = mul_tuple(mul_tuple(rest_shape, tile_shape), dilation)
+    # View, not copy
+    rest_shape_ = reversed(rest_shape) if flip_tiled_dims else rest_shape
+    tile_shape_ = reversed(tile_shape) if flip_tiled_dims else tile_shape
+    dilation_ = reversed(dilation) if flip_tiled_dims else dilation
+    tensor_tiled = tensor.view(
+        batch_actual, *dilation_, *rest_shape_, *tile_shape_, heads, dim
+    )
+    if not is_torch_compiling():
+        assert tensor_tiled.data_ptr() == tensor.data_ptr()
+    # Undo permutation
+    # batch
+    permutation_idxes = [0]
+    # dilation, rest, tile -> rest, tile, dilation
+    for i in range(na_dim):
+        if flip_tiled_dims:
+            permutation_idxes += [2 * na_dim - i, 3 * na_dim - i, na_dim - i]
+        else:
+            permutation_idxes += [na_dim + i + 1, 2 * na_dim + i + 1, i + 1]
+    # heads, head_dim
+    permutation_idxes += [na_dim * 3 + 1, na_dim * 3 + 2]
+    # View, not copy
+    tensor_permuted = tensor_tiled.permute(*permutation_idxes)
+    if not is_torch_compiling():
+        assert tensor_permuted.data_ptr() == tensor_tiled.data_ptr()
+    # Reshape back and copy
+    out = tensor_permuted.reshape(
+        batch_actual, *token_layout_padded, heads, dim
+    ).contiguous()
+    # NOTE: token permute without dilation is a no-op for 1-D
+    # assert na_dim == 1 or out.data_ptr() != tensor_permuted.data_ptr()
+    assert out.is_contiguous()
+    return out
+def _maybe_unpad(tensor: Tensor, padding: DimensionType):
+    if tensor.dim() not in [4, 5, 6]:
+        raise ValueError(
+            "Expected 4D, 5D, or 6D tensor (corresponding to NA1D, 2D, 3D), "
+            f"got {tensor.dim()}D input."
+        )
+    na_dim = tensor.dim() - 3
+    assert na_dim in [1, 2, 3]
+    if len(padding) != na_dim:
+        raise ValueError(
+            f"Expected {na_dim}D padding shape for NA{na_dim}D, " f"got {padding=}."
+        )
+    token_layout = tensor.shape[1 : na_dim + 1]
+    # Slice
+    if any(p for p in padding):
+        assert all(p >= 0 for p in padding)
+        orig_lens = tuple(x - p for x, p in zip(token_layout, padding))
+        # TODO: there must be a better way
+        if len(orig_lens) == 1:
+            x = orig_lens[0]
+            return tensor[:, :x].contiguous()
+        elif len(orig_lens) == 2:
+            x, y = orig_lens
+            return tensor[:, :x, :y].contiguous()
+        elif len(orig_lens) == 3:
+            x, y, z = orig_lens
+            return tensor[:, :x, :y, :z].contiguous()
+        else:
+            raise NotImplementedError()
+    return tensor
+def token_permute_torch(
+    tensor: Tensor,
+    tile_shape: DimensionType,
+    dilation: DimensionType,
+    flip_tiled_dims: bool,
+) -> Tensor:
+    if tensor.dim() not in [4, 5, 6]:
+        raise ValueError(
+            "Expected 4D, 5D, or 6D tensor (corresponding to NA1D, 2D, 3D), "
+            f"got {tensor.dim()}D input."
+        )
+    tensor_pad = _maybe_pad(tensor, tile_shape=tile_shape, dilation=dilation)
+    output = _token_permute(
+        tensor_pad,
+        tile_shape=tile_shape,
+        dilation=dilation,
+        flip_tiled_dims=flip_tiled_dims,
+    )
+    return output
+def token_unpermute_torch(
+    tensor: Tensor,
+    token_layout: DimensionType,
+    tile_shape: DimensionType,
+    dilation: DimensionType,
+    flip_tiled_dims: bool,
+) -> Tensor:
+    if tensor.dim() != 4:
+        raise ValueError(f"Expected flattened 4D tensor, got {tensor.dim()}D input.")
+    token_layout_padded = mul_tuple(
+        mul_tuple(
+            ceil_div_tuple(ceil_div_tuple(token_layout, tile_shape), dilation),
+            dilation,
+        ),
+        tile_shape,
+    )
+    padding = sub_tuple(token_layout_padded, token_layout)
+    output = _maybe_unpad(
+        _token_unpermute(
+            tensor,
+            token_layout=token_layout,
+            tile_shape=tile_shape,
+            dilation=dilation,
+            flip_tiled_dims=flip_tiled_dims,
+        ),
+        padding=padding,
+    )
+    return output
+__all__ = [
+    "token_permute_torch",
+    "token_unpermute_torch",
+]

build/torch210-cu126-aarch64-linux/utils/__init__.py ADDED Viewed

	@@ -0,0 +1,22 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################

build/torch210-cu126-aarch64-linux/utils/checks.py ADDED Viewed

	@@ -0,0 +1,726 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+import functools
+from collections.abc import Sequence
+from typing import Any, Optional, Tuple, Union
+import torch  # noqa: F401
+from torch import Tensor
+from .._types import CausalArgType, DimensionType, KernelSchedule, NoneType
+from ..utils import log
+from ..utils.environment import is_torch_compiling
+from ..utils.tuples import create_causal_arg_from_bool, create_dim_from_int
+from ..utils.varlen import generate_varlen_parameters
+logger = log.get_logger(__name__)
+def log_or_raise_error(
+    msg: str, raise_error: bool = False, exception: Any = RuntimeError
+):
+    if raise_error:
+        raise exception(msg)
+    else:
+        logger.debug(msg)
+def _universal_tensor_checks(
+    query: Tensor, key: Tensor, value: Tensor, raise_error: bool = True
+) -> bool:
+    target_fn = functools.partial(log_or_raise_error, raise_error=raise_error)
+    if query.is_sparse or key.is_sparse or value.is_sparse:
+        target_fn(
+            "NATTEN does not support sparse tensors.", exception=NotImplementedError
+        )
+        return False
+    if query.is_nested or key.is_nested or value.is_nested:
+        target_fn(
+            "NATTEN does not support nested tensors.", exception=NotImplementedError
+        )
+        return False
+    if query.device != key.device or query.device != value.device:
+        target_fn(
+            "Query, key, and value must be on the same device, "
+            f"got {query.device=}, {key.device=}, {value.device=}.",
+            exception=ValueError,
+        )
+        return False
+    if query.dtype != key.dtype or query.dtype != value.dtype:
+        target_fn(
+            "Query, key, and value must assume the same data type, "
+            f"got {query.dtype=}, {key.dtype=}, {value.dtype=}.",
+            exception=ValueError,
+        )
+        return False
+    return True
+def na_tensor_checks(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    must_match_head_dims: bool = False,
+    supports_gqa_mqa: bool = True,
+    raise_error: bool = True,
+    backend_name: Optional[str] = None,
+) -> bool:
+    backend_name = backend_name or "This operation/backend"
+    if not _universal_tensor_checks(query, key, value):
+        return False
+    target_fn = functools.partial(log_or_raise_error, raise_error=raise_error)
+    if query.dim() != key.dim() or query.dim() != value.dim():
+        target_fn(
+            "Query, key, and value must have the same rank, "
+            f"got {query.dim()=}, {key.dim()=}, {value.dim()=}.",
+            exception=ValueError,
+        )
+        return False
+    if query.dim() not in [4, 5, 6]:
+        target_fn(
+            "Expected 4-D, 5-D, or 6-D tensors as inputs (corresponding to NA1D, NA2D, and NA3D), "
+            f"got {query.dim()=}.",
+            exception=ValueError,
+        )
+        return False
+    na_dim = query.dim() - 3  # minus batch, heads, head_dim
+    if query.shape[-1] != key.shape[-1]:
+        target_fn(
+            f"Q and K head dims must match, got {query.shape[-1]=}, {key.shape[-1]=}.",
+            exception=ValueError,
+        )
+        return False
+    if query.shape[0] != key.shape[0] or query.shape[0] != value.shape[0]:
+        target_fn(
+            "Q, K, and V must match in batch size, got "
+            f"{query.shape[0]=}, {key.shape[0]=}, {value.shape[0]=}.",
+            exception=ValueError,
+        )
+        return False
+    if must_match_head_dims and query.shape[-1] != value.shape[-1]:
+        target_fn(
+            f"{backend_name} does not support different head dims for QK and V, got "
+            f"{query.shape[-1]=}, {value.shape[-1]=}.",
+            exception=ValueError,
+        )
+        return False
+    if (
+        query.shape[1 : na_dim + 1] != key.shape[1 : na_dim + 1]
+        or query.shape[1 : na_dim + 1] != value.shape[1 : na_dim + 1]
+    ):
+        target_fn(
+            "Neighborhood Attention operations require Q, K, and V to match in their token layouts, got "
+            f"{query.shape[1:na_dim+1]=}, {key.shape[1:na_dim+1]=}, {value.shape[1:na_dim+1]=}.",
+            exception=ValueError,
+        )
+        return False
+    if not supports_gqa_mqa and (
+        query.shape[-2] != key.shape[-2] or query.shape[-2] != value.shape[-2]
+    ):
+        target_fn(
+            f"{backend_name} does not support GQA/MQA, therefore number of heads in Q, K, and V "
+            f"must match, got {query.shape[-2]=}, {key.shape[-2]=}, {value.shape[-2]=}.",
+            exception=ValueError,
+        )
+        return False
+    if supports_gqa_mqa:
+        if key.shape[-2] != value.shape[-2]:
+            target_fn(
+                "Key and value must always have the same number of heads, got "
+                f"{key.shape[-2]=}, {value.shape[-2]=}.",
+                exception=ValueError,
+            )
+            return False
+        heads_q = query.shape[-2]
+        heads_kv = key.shape[-2]
+        if heads_q < heads_kv or heads_q % heads_kv != 0:
+            target_fn(
+                "Key/value heads must evenly divide query heads, got "
+                f"{heads_q=}, {heads_kv=}.",
+                exception=ValueError,
+            )
+            return False
+    return True
+def fmha_tensor_checks(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    must_match_head_dims: bool = False,
+    supports_gqa_mqa: bool = True,
+    raise_error: bool = True,
+    backend_name: Optional[str] = None,
+) -> bool:
+    backend_name = backend_name or "This operation/backend"
+    if not _universal_tensor_checks(query, key, value):
+        return False
+    target_fn = functools.partial(log_or_raise_error, raise_error=raise_error)
+    if query.dim() != key.dim() or query.dim() != value.dim():
+        target_fn(
+            "Query, key, and value must have the same rank, "
+            f"got {query.dim()=}, {key.dim()=}, {value.dim()=}.",
+            exception=ValueError,
+        )
+        return False
+    if query.dim() != 4:
+        target_fn(
+            "Expected 4-D tensors as inputs to FMHA, " f"got {query.dim()=}.",
+            exception=ValueError,
+        )
+        return False
+    if query.shape[-1] != key.shape[-1]:
+        target_fn(
+            f"Q and K head dims must match, got {query.shape[-1]=}, {key.shape[-1]=}.",
+            exception=ValueError,
+        )
+        return False
+    if must_match_head_dims and query.shape[-1] != value.shape[-1]:
+        target_fn(
+            f"{backend_name} does not support different head dims for QK and V, got "
+            f"{query.shape[-1]=}, {value.shape[-1]=}.",
+            exception=ValueError,
+        )
+        return False
+    if query.shape[0] != key.shape[0] or query.shape[0] != value.shape[0]:
+        target_fn(
+            "Q, K, and V must match in batch size, got "
+            f"{query.shape[0]=}, {key.shape[0]=}, {value.shape[0]=}.",
+            exception=ValueError,
+        )
+        return False
+    if key.shape[1] != value.shape[1]:
+        target_fn(
+            f"K and V must match in sequence length, got {key.shape[1]=}, {value.shape[1]=}.",
+            exception=ValueError,
+        )
+        return False
+    if not supports_gqa_mqa and (
+        query.shape[-2] != key.shape[-2] or query.shape[-2] != value.shape[-2]
+    ):
+        target_fn(
+            f"{backend_name} does not support GQA/MQA, therefore number of heads in Q, K, and V "
+            f"must match, got {query.shape[-2]=}, {key.shape[-2]=}, {value.shape[-2]=}.",
+            exception=ValueError,
+        )
+        return False
+    if supports_gqa_mqa:
+        if key.shape[-2] != value.shape[-2]:
+            target_fn(
+                "Key and value must always have the same number of heads, got "
+                f"{key.shape[-2]=}, {value.shape[-2]=}.",
+                exception=ValueError,
+            )
+            return False
+        heads_q = query.shape[-2]
+        heads_kv = key.shape[-2]
+        if heads_q < heads_kv or heads_q % heads_kv != 0:
+            target_fn(
+                "Key/value heads must evenly divide query heads, got "
+                f"{heads_q=}, {heads_kv=}.",
+                exception=ValueError,
+            )
+            return False
+    return True
+def additional_kv_tensor_checks(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    add_key: Optional[Tensor] = None,
+    add_value: Optional[Tensor] = None,
+    must_match_head_dims: bool = False,
+    supports_gqa_mqa: bool = True,
+):
+    if (add_key is not None) ^ (add_value is not None):
+        raise ValueError(
+            "`additional_keys` and `additional_values` must be either both Tensors or None."
+        )
+    if add_key is None:
+        return
+    assert add_key is not None and add_value is not None
+    _universal_tensor_checks(query, add_key, add_value)
+    if query.shape[-1] != add_key.shape[-1]:
+        raise ValueError(
+            f"Q and K head dims must match, got {query.shape[-1]=}, {add_key.shape[-1]=}."
+        )
+    if must_match_head_dims and query.shape[-1] != add_value.shape[-1]:
+        raise ValueError(
+            "This operation does not support different head dims for QK and V, got "
+            f"{query.shape[-1]=}, {add_value.shape[-1]=}."
+        )
+    if query.shape[0] != add_key.shape[0] or query.shape[0] != add_value.shape[0]:
+        raise ValueError(
+            "Q, additional K, and additional V must match in batch size, got "
+            f"{query.shape[0]=}, {add_key.shape[0]=}, {add_value.shape[0]=}."
+        )
+    if add_key.shape[1] != add_value.shape[1]:
+        raise ValueError(
+            f"Additional K and V must match in sequence length, got {add_key.shape[1]=}, "
+            f"{add_value.shape[1]=}."
+        )
+    if key.shape[0] != add_key.shape[0] or value.shape[0] != add_value.shape[0]:
+        raise ValueError(
+            "Additional key/value tokens must match the self attention key/value tokens in batch "
+            f"size, got {key.shape[0]=} != {add_key.shape[0]=}, and "
+            f"{value.shape[0]=} != {add_value.shape[0]=}."
+        )
+    if key.shape[-2] != add_key.shape[-2] or value.shape[-2] != add_value.shape[-2]:
+        raise ValueError(
+            "Additional key/value tokens must match the self attention key/value tokens in number "
+            f"of heads, got {key.shape[-2]=} != {add_key.shape[-2]=}, and "
+            f"{value.shape[-2]=} != {add_value.shape[-2]=}."
+        )
+    if key.shape[-1] != add_key.shape[-1] or value.shape[-1] != add_value.shape[-1]:
+        raise ValueError(
+            "Additional key/value tokens must match the self attention key/value tokens in head "
+            f"dim, got {key.shape[-1]=} != {add_key.shape[-1]=}, and "
+            f"{value.shape[-1]=} != {add_value.shape[-1]=}."
+        )
+    if not supports_gqa_mqa and (
+        query.shape[-2] != add_key.shape[-2] or query.shape[-2] != add_value.shape[-2]
+    ):
+        raise ValueError(
+            f"This operation does not support GQA/MQA, therefore number of heads in Q, K, and V "
+            f"must match, got {query.shape[-2]=}, {key.shape[-2]=}, {value.shape[-2]=}."
+        )
+    if supports_gqa_mqa:
+        if (
+            key.shape[-2] != value.shape[-2]
+            or key.shape[-2] != add_key.shape[-2]
+            or key.shape[-2] != add_value.shape[-2]
+        ):
+            raise ValueError(
+                "Key and value, original and additional, must always have the same number of heads, got "
+                f"{key.shape[-2]=}, {value.shape[-2]=}, {add_key.shape[-2]=}, {add_value.shape[-2]=}."
+            )
+        heads_q = query.shape[-2]
+        heads_kv = key.shape[-2]
+        if heads_q < heads_kv or heads_q % heads_kv != 0:
+            raise ValueError(
+                "Key/value heads must evenly divide query heads, got "
+                f"{heads_q=}, {heads_kv=}."
+            )
+def check_input_size_arg(na_dim: int, input_size: Any) -> DimensionType:
+    assert na_dim > 0 and na_dim < 4
+    if (
+        isinstance(input_size, Sequence)
+        and len(input_size) == na_dim
+        and all(isinstance(x, int) and x > 1 for x in input_size)
+    ):
+        return tuple(x for x in input_size)
+    if isinstance(input_size, int) and input_size > 1:
+        return create_dim_from_int(na_dim, value=input_size)
+    raise ValueError(
+        "Invalid value for `input_size`; expected an integer or iterable of integers, all >= 2, "
+        f"got {type(input_size)=}, {input_size=}."
+    )
+def check_kernel_size_arg(na_dim: int, kernel_size: Any) -> DimensionType:
+    assert na_dim > 0 and na_dim < 4
+    if (
+        isinstance(kernel_size, Sequence)
+        and len(kernel_size) == na_dim
+        and all(isinstance(x, int) and x > 1 for x in kernel_size)
+    ):
+        return tuple(x for x in kernel_size)
+    if isinstance(kernel_size, int) and kernel_size > 1:
+        return create_dim_from_int(na_dim, value=kernel_size)
+    raise ValueError(
+        "Invalid value for `kernel_size`; expected an integer or iterable of integers, all >= 2, "
+        f"got {type(kernel_size)=}, {kernel_size=}."
+    )
+def check_stride_arg(na_dim: int, stride: Any) -> DimensionType:
+    assert na_dim > 0 and na_dim < 4
+    if stride is None:
+        return create_dim_from_int(na_dim, value=1)
+    if (
+        isinstance(stride, Sequence)
+        and len(stride) == na_dim
+        and all(isinstance(x, int) and x > 0 for x in stride)
+    ):
+        return tuple(x for x in stride)
+    if isinstance(stride, int) and stride > 0:
+        return create_dim_from_int(na_dim, value=stride)
+    raise ValueError(
+        "Invalid value for `stride`; expected an integer or tuple of positive integers, "
+        f"got {type(stride)=}, {stride=}."
+    )
+def check_dilation_arg(na_dim: int, dilation: Any) -> DimensionType:
+    assert na_dim > 0 and na_dim < 4
+    if dilation is None:
+        return create_dim_from_int(na_dim, value=1)
+    if (
+        isinstance(dilation, Sequence)
+        and len(dilation) == na_dim
+        and all(isinstance(x, int) and x > 0 for x in dilation)
+    ):
+        return tuple(x for x in dilation)
+    if isinstance(dilation, int) and dilation > 0:
+        return create_dim_from_int(na_dim, value=dilation)
+    raise ValueError(
+        "Invalid value for `dilation`; expected an integer or tuple of positive integers, "
+        f"got {type(dilation)=}, {dilation=}."
+    )
+def check_causal_arg(na_dim: int, is_causal: Any) -> CausalArgType:
+    assert na_dim > 0 and na_dim < 4
+    if is_causal is None:
+        return create_causal_arg_from_bool(na_dim, value=False)
+    if (
+        isinstance(is_causal, Sequence)
+        and len(is_causal) == na_dim
+        and all(isinstance(c, bool) for c in is_causal)
+    ):
+        return tuple(c for c in is_causal)
+    if isinstance(is_causal, bool):
+        return create_causal_arg_from_bool(na_dim, value=is_causal)
+    raise ValueError(
+        "Invalid value for `is_causal`; expected a boolean or tuple of booleans, "
+        f"got {type(is_causal)=}, {is_causal=}."
+    )
+def check_all_args(
+    na_dim: int, kernel_size: Any, stride: Any, dilation: Any, is_causal: Any
+) -> Tuple[DimensionType, DimensionType, DimensionType, CausalArgType]:
+    kernel_size_out, stride_out, dilation_out, is_causal_out = (
+        check_kernel_size_arg(na_dim, kernel_size),
+        check_stride_arg(na_dim, stride),
+        check_dilation_arg(na_dim, dilation),
+        check_causal_arg(na_dim, is_causal),
+    )
+    return kernel_size_out, stride_out, dilation_out, is_causal_out
+def check_args_against_input(
+    input_tensor: Tensor,
+    kernel_size: DimensionType,
+    stride: DimensionType,
+    dilation: DimensionType,
+    is_causal: CausalArgType,
+):
+    assert input_tensor.dim() in [4, 5, 6]
+    na_dim = input_tensor.dim() - 3
+    input_size = input_tensor.shape[1 : 1 + na_dim]
+    if any(k * d > x for x, k, d in zip(input_size, kernel_size, dilation)):
+        raise ValueError(
+            "The product of kernel size and dilation cannot be larger than input size "
+            f"along any dimension, got {input_size=} ({input_tensor.shape=}), "
+            f"{kernel_size=}, {dilation=}."
+        )
+    if any(s > k for k, s in zip(kernel_size, stride)):
+        raise ValueError(
+            "Stride cannot be larger than kernel size along any dimension, got "
+            f"{kernel_size=}, {stride=}."
+        )
+def is_self_attention(
+    input_tensor: Tensor,
+    kernel_size: DimensionType,
+    is_causal: CausalArgType,
+    has_additional_attention: bool,
+):
+    assert input_tensor.dim() in [4, 5, 6]
+    na_dim = input_tensor.dim() - 3
+    input_size = input_tensor.shape[1 : 1 + na_dim]
+    # Special case: 1-D causal with full window is equivalent to standard 1-D causal
+    # as long as there isn't any additional context (non causal)
+    if na_dim == 1 and not has_additional_attention:
+        return kernel_size[0] == input_size[0]
+    return all(k == x and not c for x, k, c in zip(input_size, kernel_size, is_causal))
+def check_tile_shape(
+    tile_shape: Any,
+) -> DimensionType:
+    if (
+        isinstance(tile_shape, Sequence)
+        and len(tile_shape) <= 3
+        and all(isinstance(x, int) for x in tile_shape)
+    ):
+        return tuple(x for x in tile_shape)
+    raise ValueError(
+        f"Unsupported value for tile shape; expected an iterable of at most 3 integers, "
+        f"got {type(tile_shape)=}, {tile_shape}."
+    )
+def check_kernel_schedule(kernel_schedule: Any) -> Optional[KernelSchedule]:
+    if kernel_schedule is None:
+        return None
+    if isinstance(kernel_schedule, KernelSchedule):
+        return kernel_schedule
+    if kernel_schedule == "non":
+        return KernelSchedule.NonPersistent
+    elif kernel_schedule == "coop":
+        return KernelSchedule.WarpSpecializedCooperative
+    elif kernel_schedule == "pp":
+        return KernelSchedule.WarpSpecializedPingpong
+    raise ValueError(
+        f"Kernel schedule {kernel_schedule} is invalid; choices are: "
+        "`non` (non-persistent), `coop` (warp-specialized cooperative), and "
+        "`pp` (warp-specialized ping-ponging)."
+    )
+# Varlen FMHA Checks
+def varlen_tensor_checks(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    seqlens_Q: Optional[Tensor] = None,
+    seqlens_KV: Optional[Tensor] = None,
+    cumulative_seqlen_Q: Optional[Tensor] = None,
+    cumulative_seqlen_KV: Optional[Tensor] = None,
+    max_seqlen_Q: Optional[int] = None,
+    max_seqlen_KV: Optional[int] = None,
+) -> Union[
+    Tuple[NoneType, NoneType, int, int],
+    Tuple[Tensor, Tensor, int, int],
+]:
+    if query.shape[0] != key.shape[0] or query.shape[0] != value.shape[0]:
+        raise ValueError(
+            "Q, K, and V must match in batch size, got "
+            f"{query.shape[0]=}, {key.shape[0]=}, {value.shape[0]=}."
+        )
+    if all(
+        x is None
+        for x in [
+            seqlens_Q,
+            seqlens_KV,
+            cumulative_seqlen_Q,
+            cumulative_seqlen_KV,
+        ]
+    ) and all(
+        x is None or x == 0
+        for x in [
+            max_seqlen_Q,
+            max_seqlen_KV,
+        ]
+    ):
+        # Not varlen
+        return None, None, 0, 0
+    if seqlens_Q is not None or seqlens_KV is not None:
+        # Generate cumulative_seqlen_{Q,KV}, max_seqlen_{Q,KV}, total_seqlen_{Q,KV}
+        # based on user input
+        return generate_varlen_parameters(
+            query=query,
+            key=key,
+            value=value,
+            seqlens_Q=seqlens_Q,
+            seqlens_KV=seqlens_KV,
+        )
+    # Validate user-input cumulative_seqlen_{Q,KV}, max_seqlen_{Q,KV}, total_seqlen_{Q,KV}
+    if any(
+        x is None
+        for x in [
+            cumulative_seqlen_Q,
+            cumulative_seqlen_KV,
+            max_seqlen_Q,
+            max_seqlen_KV,
+        ]
+    ):
+        raise ValueError(
+            "Variable length Attention requires all of "
+            "cumulative_seqlen_{Q,KV} and max_seqlen_{Q,KV} to be set."
+        )
+    if query.shape[0] != 1:
+        raise ValueError(
+            "Variable length Attention only supports sequence-packed memory layout "
+            f"(batch = 1), got {query.shape[0]=}."
+        )
+    assert cumulative_seqlen_Q is not None
+    assert cumulative_seqlen_KV is not None
+    assert max_seqlen_Q is not None
+    assert max_seqlen_KV is not None
+    if not isinstance(max_seqlen_Q, int) or not isinstance(max_seqlen_KV, int):
+        raise ValueError(
+            "max_seqlen_Q and max_seqlen_KV must be ints, got "
+            f"{type(max_seqlen_Q)=}, {type(max_seqlen_KV)=}, {max_seqlen_Q=}, {max_seqlen_KV=}."
+        )
+    total_seqlen_Q = query.shape[1]
+    total_seqlen_KV = key.shape[1]
+    if max_seqlen_Q > total_seqlen_Q:
+        raise ValueError(
+            "Maximum sequence length cannot exceed total, got "
+            f"{max_seqlen_Q=}, {total_seqlen_Q=}."
+        )
+    if max_seqlen_KV > total_seqlen_KV:
+        raise ValueError(
+            "Maximum sequence length cannot exceed total, got "
+            f"{max_seqlen_KV=}, {total_seqlen_KV=}."
+        )
+    # NOTE: this check introduces recompiles
+    if not is_torch_compiling():
+        if (max_seqlen_Q == 0) != (max_seqlen_KV == 0):
+            raise ValueError(
+                "max_seqlen_Q and max_seqlen_KV must both be zero or both be non-zero, got "
+                f"{max_seqlen_Q=}, {max_seqlen_KV=}."
+            )
+    if max_seqlen_Q < 0 or max_seqlen_KV < 0:
+        raise ValueError(
+            "Maximum sequence length cannot be negative, got "
+            f"{max_seqlen_Q=}, {max_seqlen_KV=}."
+        )
+    if not isinstance(cumulative_seqlen_Q, Tensor) or not isinstance(
+        cumulative_seqlen_KV, Tensor
+    ):
+        raise ValueError(
+            "cumulative_seqlen_Q and cumulative_seqlen_KV must both be tensors."
+        )
+    if (
+        cumulative_seqlen_Q.device != query.device
+        or cumulative_seqlen_KV.device != query.device
+    ):
+        raise ValueError(
+            "cumulative_seqlen_Q and cumulative_seqlen_KV must be on the same device as QKV, but "
+            f"{cumulative_seqlen_Q.device=}, {cumulative_seqlen_KV.device=}, {query.device=}."
+        )
+    if (
+        cumulative_seqlen_Q.dtype != torch.int32
+        or cumulative_seqlen_KV.dtype != torch.int32
+    ):
+        raise ValueError(
+            "cumulative_seqlen_Q and cumulative_seqlen_KV must both be torch.int32 tensors, got "
+            f"{cumulative_seqlen_Q.dtype=}, {cumulative_seqlen_KV.dtype=}."
+        )
+    if cumulative_seqlen_Q.dim() != 1 or cumulative_seqlen_KV.dim() != 1:
+        raise ValueError(
+            "cumulative_seqlen_Q and cumulative_seqlen_KV must both be 1-D tensors, got "
+            f"{cumulative_seqlen_Q.dim()=}, {cumulative_seqlen_KV.dim()=}."
+        )
+    if cumulative_seqlen_Q.shape[0] != cumulative_seqlen_KV.shape[0]:
+        raise ValueError(
+            "cumulative_seqlen_Q and cumulative_seqlen_KV must match in size, got "
+            f"{cumulative_seqlen_Q.shape=}, {cumulative_seqlen_KV.shape=}."
+        )
+    if cumulative_seqlen_Q.shape[0] < 2:
+        raise ValueError(
+            "cumulative_seqlen_Q and cumulative_seqlen_KV must contain at least 2 elements, got "
+            f"{cumulative_seqlen_Q.shape=}, {cumulative_seqlen_KV.shape=}."
+        )
+    return (
+        cumulative_seqlen_Q,
+        cumulative_seqlen_KV,
+        max_seqlen_Q,
+        max_seqlen_KV,
+    )

build/torch210-cu126-aarch64-linux/utils/device.py ADDED Viewed

	@@ -0,0 +1,50 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+from typing import Optional
+import torch
+def is_cuda(device: torch.device) -> bool:
+    return torch.cuda.is_available() and torch.version.cuda and device.type == "cuda"  # type: ignore
+def is_rocm(device: torch.device) -> bool:
+    return torch.cuda.is_available() and torch.version.hip and device.type == "cuda"  # type: ignore
+def is_cpu(device: torch.device) -> bool:
+    return device.type == "cpu"
+def get_device_cc(device: Optional[torch.device] = None) -> int:
+    if (
+        torch.cuda.is_available()
+        and torch.version.cuda
+        and (device is None or is_cuda(device))
+    ):
+        major, minor = torch.cuda.get_device_capability(device)
+        return major * 10 + minor
+    return 0

build/torch210-cu126-aarch64-linux/utils/dtype.py ADDED Viewed

	@@ -0,0 +1,36 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+import torch
+def is_full(dtype: torch.dtype) -> bool:
+    return dtype == torch.float32
+def is_half(dtype: torch.dtype) -> bool:
+    return dtype in [torch.float16, torch.bfloat16]
+def is_fp8(dtype: torch.dtype) -> bool:
+    return dtype in [torch.float8_e5m2, torch.float8_e4m3fn]

build/torch210-cu126-aarch64-linux/utils/environment.py ADDED Viewed

	@@ -0,0 +1,79 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+import os
+import torch
+from ..utils.device import get_device_cc
+def parse_env_flag(env_var: str, default: bool) -> bool:
+    default_str = "1" if default else "0"
+    out_str = os.getenv(env_var, default_str)
+    if out_str.strip() == "":
+        return default
+    if out_str == "0":
+        return False
+    if out_str == "1":
+        return True
+    return default
+def parse_env_int(env_var: str, default: int) -> int:
+    out_str = os.getenv(env_var, str(default))
+    if out_str.strip() == "":
+        return default
+    try:
+        return int(out_str)
+    except ValueError:
+        return default
+def parse_env_str(env_var: str, default: str) -> str:
+    return os.getenv(env_var, str(default))
+_IS_CUDA_AVAILABLE = torch.cuda.is_available()
+_TORCH_VERSION = [int(x) for x in torch.__version__.split(".")[:2]]
+_IS_TORCH_COMPILE_SUPPORTED = _TORCH_VERSION >= [2, 6] and get_device_cc() >= 70
+# Guard registering libnatten APIs as torch ops with environment variables
+# In case any unusual bugs from torch compile come up again
+# Also restrict to torch 2.8 and later
+# https://github.com/pytorch/pytorch/issues/137979#issuecomment-3614956989
+DISABLE_TORCH_OPS = _TORCH_VERSION < [2, 8] or parse_env_flag(
+    "NATTEN_DISABLE_TORCH_OPS", False
+)
+# Controls all regions guarded against torch compile
+# Logs, and certain assertions cause graph breaks.
+def is_torch_compiling() -> bool:
+    try:
+        return torch.compiler.is_compiling()
+    except:
+        # Assume too old to support torch compile
+        return False

build/torch210-cu126-aarch64-linux/utils/log.py ADDED Viewed

	@@ -0,0 +1,134 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+import enum
+import logging
+import os
+import sys
+from ..utils.environment import is_torch_compiling, parse_env_str
+log_format = "| %(asctime)s | [[ %(name)s ]] [ %(levelname)s ]: %(message)s"
+class LogLevel(enum.Enum):
+    Default = 0
+    Debug = 1
+    Info = 2
+    Warnings = 3
+    Errors = 4
+    Critical = 5
+def _get_log_level() -> LogLevel:
+    log_level = parse_env_str("NATTEN_LOG_LEVEL", "").lower()
+    if log_level == "debug":
+        return LogLevel.Debug
+    elif log_level == "info":
+        return LogLevel.Info
+    elif log_level == "warning":
+        return LogLevel.Warnings
+    elif log_level == "error":
+        return LogLevel.Errors
+    elif log_level == "critical":
+        return LogLevel.Critical
+    return LogLevel.Default
+_map_log_level = {
+    LogLevel.Default: logging.INFO,
+    LogLevel.Debug: logging.DEBUG,
+    LogLevel.Info: logging.INFO,
+    LogLevel.Warnings: logging.WARNING,
+    LogLevel.Errors: logging.ERROR,
+    LogLevel.Critical: logging.CRITICAL,
+}
+# Tests will stream into stderr instead of stdout
+# It can be set to either stderr, stdout or any writeable file.
+# Otherwise logging will be disabled.
+def _get_log_pipe():
+    log_pipe = parse_env_str("NATTEN_LOG_PIPE", "stdout")
+    # Skip checking /dev/null writablity
+    if log_pipe == "/dev/null":
+        return None
+    if log_pipe.lower() == "stderr":
+        return sys.stderr
+    if log_pipe.lower() == "stdout":
+        return sys.stdout
+    # Treat as file path; validate writability
+    if os.path.isfile(log_pipe) and os.access(log_pipe, os.W_OK):
+        return log_pipe
+    try:
+        open(log_pipe, "a").close()
+        return log_pipe
+    except OSError:
+        pass
+    return None
+class NattenLogger:
+    def __init__(self, name: str):
+        self.logger = logging.getLogger(name)
+        self.log_level = _map_log_level[_get_log_level()]
+        self.logger.setLevel(self.log_level)
+        self.formatter = logging.Formatter(log_format)
+        log_pipe = _get_log_pipe()
+        if log_pipe in [sys.stderr, sys.stdout]:
+            self.handler = logging.StreamHandler(log_pipe)
+        elif isinstance(log_pipe, str):
+            self.handler = logging.FileHandler(log_pipe)
+        else:
+            # Invalid / null
+            self.handler = logging.NullHandler()  # type: ignore[assignment]
+        self.handler.setLevel(self.log_level)
+        self.handler.setFormatter(self.formatter)
+        self.logger.addHandler(self.handler)
+    def is_safe_to_log(self) -> bool:
+        return not is_torch_compiling()
+    def info(self, *args, **kwargs):
+        if self.is_safe_to_log():
+            self.logger.info(*args, **kwargs)
+    def debug(self, *args, **kwargs):
+        if self.is_safe_to_log():
+            self.logger.debug(*args, **kwargs)
+    def warning(self, *args, **kwargs):
+        if self.is_safe_to_log():
+            self.logger.warning(*args, **kwargs)
+def get_logger(name) -> NattenLogger:
+    return NattenLogger(name)

build/torch210-cu126-aarch64-linux/utils/tensor.py ADDED Viewed

	@@ -0,0 +1,113 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+from typing import Optional
+import torch
+from torch import Size, Tensor
+def _get_expected_attn_shape(input_tensor: Tensor, attention_dim: int) -> Size:
+    shape = [x for x in input_tensor.shape[:-1]] + [attention_dim]
+    return Size(shape)
+def make_attn_tensor_from_input(input_tensor: Tensor, attention_dim: int) -> Tensor:
+    return torch.empty(
+        _get_expected_attn_shape(input_tensor, attention_dim),
+        device=input_tensor.device,
+        dtype=input_tensor.dtype,
+        requires_grad=input_tensor.requires_grad,
+    )
+def check_additional_keys(
+    input_tensor: Tensor, additional_keys: Optional[Tensor]
+) -> int:
+    if additional_keys is None:
+        return 0
+    if additional_keys.dim() != 4:
+        raise ValueError(
+            "Additional tokens have to be shaped as a rank-4 tensor; "
+            f"got {additional_keys.dim()}."
+        )
+    batch_size, heads, tokens, dim = additional_keys.shape
+    expected_batch_size = input_tensor.shape[0]
+    expected_heads = input_tensor.shape[1]
+    expected_dim = input_tensor.shape[-1]
+    if (
+        batch_size != expected_batch_size
+        or expected_heads != heads
+        or expected_dim != dim
+    ):
+        raise ValueError(
+            "Shape mismatch between input tensor and additional tokens; "
+            "they must match in batch size, heads, and dim per head. "
+            f"Got {input_tensor.shape=}, {additional_keys.shape=}."
+        )
+    return tokens
+def check_additional_values(
+    attn_tensor: Tensor,
+    additional_values: Optional[Tensor],
+    value: Tensor,
+    expected_attn_weights: int,
+) -> int:
+    if additional_values is None and attn_tensor.shape[-1] == expected_attn_weights:
+        return 0
+    if additional_values is None:
+        raise ValueError(
+            f"Expected {expected_attn_weights} attention weights per token, "
+            f"got {attn_tensor.shape[-1]=}."
+        )
+    if additional_values.dim() != 4:
+        raise ValueError(
+            "Additional tokens have to be shaped as a rank-4 tensor; "
+            f"got {additional_values.dim()}."
+        )
+    if additional_values.shape[-1] != value.shape[-1]:
+        raise ValueError(
+            "Additional value tokens must match the dimension of the "
+            f"rest of the tokens, got {additional_values.shape[-1]=} != "
+            f"{value.shape[-1]=}."
+        )
+    batch_size, heads, tokens, dim = additional_values.shape
+    if tokens + expected_attn_weights != attn_tensor.shape[-1]:
+        raise ValueError(
+            f"Expected {expected_attn_weights + tokens} attention weights per token, "
+            f"got {attn_tensor.shape[-1]=}."
+        )
+    expected_batch_size = attn_tensor.shape[0]
+    expected_heads = attn_tensor.shape[1]
+    if batch_size != expected_batch_size or expected_heads != heads:
+        raise ValueError(
+            "Shape mismatch between attention tensor and additional tokens; "
+            "they must match in batch size and heads. "
+            f"Got {attn_tensor.shape=}, {additional_values.shape=}."
+        )
+    return tokens

build/torch210-cu126-aarch64-linux/utils/testing.py ADDED Viewed

	@@ -0,0 +1,149 @@

+#################################################################################################
+# Copyright (c) 2022 - 2026 Ali Hassani.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################################
+import torch
+from .._environment import _IS_CUDA_AVAILABLE, _RUN_EXTENDED_TESTS, HAS_LIBNATTEN
+from ..backends.flex import _FLEX_COMPILE_SUPPORTED, _FLEX_SUPPORTED
+from ..utils.device import get_device_cc, is_cuda
+def skip_if_libnatten_is_not_supported():
+    def decorator(f):
+        def wrapper(self, *args, **kwargs):
+            if not _IS_CUDA_AVAILABLE:
+                self.skipTest("CUDA is not available.")
+            elif not HAS_LIBNATTEN:
+                self.skipTest("Libnatten is not available.")
+            else:
+                return f(self, *args, **kwargs)
+        return wrapper
+    return decorator
+def skip_if_cuda_is_not_supported():
+    def decorator(f):
+        def wrapper(self, *args, **kwargs):
+            if not _IS_CUDA_AVAILABLE:
+                self.skipTest("CUDA is not available.")
+            else:
+                return f(self, *args, **kwargs)
+        return wrapper
+    return decorator
+def skip_if_flex_is_not_supported():
+    def decorator(f):
+        def wrapper(self, *args, **kwargs):
+            if not _FLEX_SUPPORTED or get_device_cc() < 70:
+                self.skipTest("Flex backend is not supported.")
+            else:
+                return f(self, *args, **kwargs)
+        return wrapper
+    return decorator
+def skip_if_flex_compile_is_not_supported():
+    def decorator(f):
+        def wrapper(self, *args, **kwargs):
+            if not _FLEX_COMPILE_SUPPORTED:
+                self.skipTest("Flex (compiled) backend is not supported.")
+            else:
+                return f(self, *args, **kwargs)
+        return wrapper
+    return decorator
+def skip_if_not_running_extended_tests():
+    def decorator(f):
+        def wrapper(self, *args, **kwargs):
+            if not _RUN_EXTENDED_TESTS:
+                self.skipTest("Skipping extended test cases.")
+            else:
+                return f(self, *args, **kwargs)
+        return wrapper
+    return decorator
+def skip_if_hopper_kernels_not_supported():
+    def decorator(f):
+        def wrapper(self, *args, **kwargs):
+            if get_device_cc() != 90:
+                self.skipTest("Hopper kernels are only supported on SM90.")
+            else:
+                return f(self, *args, **kwargs)
+        return wrapper
+    return decorator
+def skip_if_blackwell_kernels_not_supported():
+    def decorator(f):
+        def wrapper(self, *args, **kwargs):
+            if get_device_cc() not in [100, 103]:
+                self.skipTest(
+                    "Blackwell kernels are only supported on SM100 and SM103."
+                )
+            else:
+                return f(self, *args, **kwargs)
+        return wrapper
+    return decorator
+def supports_float16(device: torch.device) -> bool:
+    if is_cuda(device):
+        device_cc = get_device_cc(device)
+        if device_cc < 50:
+            return False
+        return True
+    # TODO:
+    return True
+def supports_bfloat16(device: torch.device) -> bool:
+    if is_cuda(device):
+        device_cc = get_device_cc(device)
+        if device_cc < 80:
+            return False
+        return True
+    # TODO:
+    return False