Kernels:

kernels-staging
/

natten

Trusted publisher

kernels-bot commited on 13 days ago

Commit

e2117e8

verified ·

1 Parent(s): 586cdf4

Uploaded using `kernel-builder`.

Browse files

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

Files changed (50) hide show

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

build/torch211-cxx11-cu126-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,178 @@

+#################################################################################################
+# 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 .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__",
+    "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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-linux/_natten_cuda_3641131.abi3.so ADDED Viewed

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:a38a8d19b1065ea660148531865732a9ec754be0efb551fcaeef77ce8d71d890
+size 103606392

build/torch211-cxx11-cu126-x86_64-linux/_ops.py ADDED Viewed

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

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

build/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,81 @@

+{
+  "name": "natten",
+  "id": "_natten_cuda_3641131",
+  "version": 1,
+  "license": "MIT",
+  "upstream": "https://github.com/SHI-Labs/NATTEN",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "8.0",
+      "9.0",
+      "9.0a"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "gum9e68BeqHOe1E0Pa6VqIUe18hV1/LH0hbn3J4ytj8=",
+      "_environment.py": "9P189XTk2YMyWuRNI5eCI3UvHRqVR8Zh2hgJsHP7MYI=",
+      "_libnatten/__init__.py": "yMW+2Kd5Nxs716jaNq15IleYQ86HR5u7SMEQDqVVunI=",
+      "_libnatten/torch_wrappers.py": "v7vkrytkghNz+cJI7l2PXGqWactsSM1ImwCsHfyQ3iA=",
+      "_natten_cuda_3641131.abi3.so": "o4qNGbEGXqZgFIUxhlcyqex1S+DvtVH8ru93zo1x2JA=",
+      "_ops.py": "KdQwUQbokH7RjoPK3Y6hQVR6zOhKfYdEPoQ1qvrIRms=",
+      "_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=",
+      "natten/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
+      "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.17.0-dev0",
+      "sha": "ce92bd77c807ce3a2f4b85d8bab69caf027b51d0",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "3641131fa0a08b8174b3035cb79e5cf27bd5baad",
+      "dirty": false
+    }
+  }
+}

build/torch211-cxx11-cu126-x86_64-linux/metadata.json.sigstore ADDED Viewed

	@@ -0,0 +1 @@

+ {"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSjCCBtGgAwIBAgIUENWDkPegFVvZRJM6q8qMhUuMwLYwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNzI5MDkyMTM0WhcNMjYwNzI5MDkzMTM0WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEdcq8BxBrq06UzALZ5aKGyV1ze0u7pb0jLRBda6h/713y6d9xSP9cMdnfrnlQjnbleRt+Ok7NLTIgUAy7rXutS6OCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQU64vuL4RzrTPjhn01iVDcqNT9rHwwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoMjBhNjEzYzEwZDliMWM5Y2I0M2E1MWU4NTBlOWZhMDQ2OWE5NzE5OTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoMjBhNjEzYzEwZDliMWM5Y2I0M2E1MWU4NTBlOWZhMDQ2OWE5NzE5OTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoMjBhNjEzYzEwZDliMWM5Y2I0M2E1MWU4NTBlOWZhMDQ2OWE5NzE5OTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDIwYTYxM2MxMGQ5YjFjOWNiNDNhNTFlODUwZTlmYTA0NjlhOTcxOTkwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzA0MTgyNTI4NjkvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABn60t5mgAAAQDAEcwRQIgLUvkKk7oLX3B1UWy0X+Ch8UaFJEDl2EeSAEXM5+1a1oCIQDOWS4Jgt9qjKDL1aAqGDgONFF3WMIm/o/GJtWVFQ+snjAKBggqhkjOPQQDAwNnADBkAjBTqD6LtwE6ptoxeZRDivuw5uK1fsfPjFpmohRPD/euDT12EwjBMk7iRplRIGBseh8CMHu6bQRud0BCGIl7QJ55GAmRZmXYTQ8Y5VZTg7IyK7rPG4I9aaMiU9rzv7skJMxh/w=="}, "tlogEntries":[{"logIndex":"2280149228", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1785316894", "inclusionPromise":{"signedEntryTimestamp":"MEUCIG1ry/M5a8Q+Zyp4icfw0vKvjTOu5i8SQyztovlzxmVdAiEAwHju/lVvwJ3JkEkdWiTpuOkt2r0lPghXA1VFjFlIw28="}, "inclusionProof":{"logIndex":"2158244966", "rootHash":"bP7TpLzXp35jXCgw2q1PbA0L9n++o7lG9fyD6jOXIws=", "treeSize":"2158244969", "hashes":["VleMKHTzOxE/Vx4/PHWMo+XAmrpe5Blq6nTIjr0HVTA=", "cg+uSPas8GSp/Xpk1dyZMoM2AY4AWFBpNQ/+mUY91k0=", "swwiNtY8cPozMCjx7lEmyW50d36M1Kz23Nu5dDd+sZQ=", "qaYLnSshaNGiEaRlCETE/NaWT6ZItJQ4IOINv1ad63E=", "8CZ1EqhgyxGzJym/Y5ujtMUP4B7JUw/hSYjhV6H7YHM=", "sS+fl5SKwsQjQE6HrC426ByW+1/o21xz4dSeBr22cqY=", "DwNueYXiyuVjob7c8CRNRN57ioovkJTOeOeOQIi6/pI=", "bXYdybsl/2BK7Yled3FQJEE3Hk0TeGcWnXUxOkVZv8w=", "xrwx0yHkH5ZmCgeWs/jLZZ8RcejqpeuUeAWNJkAYrHk=", "8MN6j2GHUjPnNtIwpas9l/XjdwE0/KPUhFDlrVkngko=", "b+xUZfuENQxvSOJxzNvYvRG8eVphfszPpZmuf4/cQ6c=", "OVsvZCKnWA+498QUIaQCtitUT6huDbC7SmhH1l8MxXI=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2158244969\nbP7TpLzXp35jXCgw2q1PbA0L9n++o7lG9fyD6jOXIws=\n\n— rekor.sigstore.dev wNI9ajBGAiEAnh/bQrjUVU/DKYSDXrcfenN3Nvhk+jyyz6tJiqR6VCECIQC4O/QWRaJOz5zc14uWgveqlFIBfe0saoa77r7MayGBlA==\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI3ZWNhN2Y5Nzg4MDhjNzJiNDAxOGJkN2VhYWMzYzVjNDAxYWRlNDZjMTdmYjEyYWU5NzAyNjM4NWE3MmExYzQ2In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FUUNJRisrKy9DZzVHMTRqUGFVRGRIbUtDb1lZVW8rZVdhc3JkYkdodDczNVRnNUFpQmpqZ2UrenlDUGRjeXFLLzMrZS92bzBqVFpPQlE2ODNBZzhWVGhtdWZ2bVE9PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRha05EUW5SSFowRjNTVUpCWjBsVlJVNVhSR3RRWldkR1ZuWmFVa3BOTm5FNGNVMW9WWFZOZDB4WmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDU2U1RWTlJHdDVUVlJOTUZkb1kwNU5hbGwzVG5wSk5VMUVhM3BOVkUwd1YycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZrWTNFNFFuaENjbkV3TmxWNlFVeGFOV0ZMUjNsV01YcGxNSFUzY0dJd2FreFNRbVFLWVRab0x6Y3hNM2syWkRsNFUxQTVZMDFrYm1aeWJteFJhbTVpYkdWU2RDdFBhemRPVEZSSloxVkJlVGR5V0hWMFV6WlBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlUyTkhaMUNrdzBVbnB5VkZCcWFHNHdNV2xXUkdOeFRsUTVja2gzZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU5ha0pvVG1wRmVsbDZSWGRhUkd4cFRWZE5OVmt5U1RCTk1rVXhDazFYVlRST1ZFSnNUMWRhYUUxRVVUSlBWMFUxVG5wRk5VOVVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDFxUW1oT2FrVjZXWHBGZDFwRWJHbE5WMDAxV1RKSk1FMHlSVEZOVjFVMFRsUkNiRTlYV21nS1RVUlJNazlYUlRWT2VrVTFUMVJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDFxUW1oT2FrVjZXWHBGZDFwRWJHbE5WMDAxV1RKSk1FMHlSVEVLVFZkVk5FNVVRbXhQVjFwb1RVUlJNazlYUlRWT2VrVTFUMVJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUkpkMWxVV1hnS1RUSk5lRTFIVVRWWmFrWnFUMWRPYVU1RVRtaE9WRVpzVDBSVmQxcFViRzFaVkVFd1RtcHNhRTlVWTNoUFZHdDNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWtFd1RWUm5lVTVVU1RST2FtdDJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNDJNSFExYldkQlFVRlJSQXBCUldOM1VsRkpaMHhWZG10TGF6ZHZURmd6UWpGVlYza3dXQ3REYURoVllVWktSVVJzTWtWbFUwRkZXRTAxS3pGaE1XOURTVkZFVDFkVE5FcG5kRGx4Q21wTFJFd3hZVUZ4UjBSblQwNUdSak5YVFVsdEwyOHZSMHAwVjFaR1VTdHpibXBCUzBKblozRm9hMnBQVUZGUlJFRjNUbTVCUkVKclFXcENWSEZFTmt3S2RIZEZObkIwYjNobFdsSkVhWFoxZHpWMVN6Rm1jMlpRYWtad2JXOW9VbEJFTDJWMVJGUXhNa1YzYWtKTmF6ZHBVbkJzVWtsSFFuTmxhRGhEVFVoMU5ncGlVVkoxWkRCQ1EwZEpiRGRSU2pVMVIwRnRVbHB0V0ZsVVVUaFpOVlphVkdjM1NYbExOM0pRUnpSSk9XRmhUV2xWT1hKNmRqZHphMHBOZUdndmR6MDlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgTRh2Eax7lxF4e3pxV97waQT6rugjVfOy5zU6VruHlNoCFQCzP0L3DFHjsreAIT7slbSE22uoJhgPMjAyNjA3MjkwOTIxMzRaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHbMIIB1wIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDcyOTA5MjEzNFowLwYJKoZIhvcNAQkEMSIEIJvXfubflBPfSzp3U8kYq+xQV4Aq0cKj0OZVocVa4XgrMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRnMGUCMQDyIRWFENKc/oHwIdQh1+VD+l7q51BeMBuIvHCOkw9M/7GDIDs0BI1Eb2NrIhWdEvwCMCIC9LZ/FUVL0d3vEMgfWM1BmJ9SgYP4s4HuqRo2C4qgeWH81djcTRO3BzCyS3J7yA=="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"fsp/l4gIxytAGL1+qsPFxAGt5GwX+xKulwJjhacqHEY="}, "signature":"MEQCIF+++/Cg5G14jPaUDdHmKCoYYUo+eWasrdbGht735Tg5AiBjjge+zyCPdcyqK/3+e/vo0jTZOBQ683Ag8VThmufvmQ=="}}

build/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-linux/natten/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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/torch211-cxx11-cu126-x86_64-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

build/torch211-cxx11-cu126-x86_64-linux/utils/tuples.py ADDED Viewed

	@@ -0,0 +1,51 @@

+#################################################################################################
+# 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 .._types import CausalArgType, DimensionType
+def ceil_div_int(x: int, y: int) -> int:
+    return (x + y - 1) // y
+def ceil_div_tuple(X: tuple, Y: tuple) -> tuple:
+    assert len(X) == len(Y)
+    return tuple(ceil_div_int(x, y) for x, y in zip(X, Y))
+def mul_tuple(X: tuple, Y: tuple) -> tuple:
+    assert len(X) == len(Y)
+    return tuple(x * y for x, y in zip(X, Y))
+def sub_tuple(X: tuple, Y: tuple) -> tuple:
+    assert len(X) == len(Y)
+    return tuple(x - y for x, y in zip(X, Y))
+def create_dim_from_int(na_dim: int, value: int) -> DimensionType:
+    return tuple(value for _ in range(na_dim))  # type: ignore
+def create_causal_arg_from_bool(na_dim: int, value: bool) -> CausalArgType:
+    return tuple(value for _ in range(na_dim))  # type: ignore

build/torch211-cxx11-cu126-x86_64-linux/utils/varlen.py ADDED Viewed

	@@ -0,0 +1,135 @@

+#################################################################################################
+# 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, Tuple, Union
+import torch  # noqa: F401
+from torch import Tensor
+from .._types import NoneType
+from ..utils.environment import is_torch_compiling
+def generate_varlen_parameters(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    seqlens_Q: Optional[Tensor] = None,
+    seqlens_KV: Optional[Tensor] = None,
+) -> Union[
+    Tuple[NoneType, NoneType, int, int],
+    Tuple[Tensor, Tensor, int, int],
+]:
+    # NOTE: max_seqlen_{Q,KV} require a device-host sync, since they're expected to be ints (with
+    # which we launch the varlen kernel) and not device tensors.
+    # .item() introduces control flow and breaks the graph.
+    # It is also inefficient to repeat this per-op, and mostly there for convenience.
+    # generate_varlen_parameters should ideally always be called by the user ahead of model
+    # forward / backward.
+    if is_torch_compiling():
+        raise RuntimeError(
+            "Running 'generate_varlen_parameters' in a torch-compiled region is disallowed as it "
+            "results in graph breaks. Please consider calling ahead of time and pass "
+            "'cumulative_seqlen_{Q,KV}' and 'max_seqlen_{Q,KV}' instead of 'seqlens_{Q,KV}' to "
+            "'attention'. "
+        )
+    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 (seqlens_Q is None) ^ (seqlens_KV is None):
+        raise ValueError(
+            "Variable length Attention requires both of seqlens_Q and seqlens_KV to be set, got "
+            f"{seqlens_Q=}, {seqlens_KV=}."
+        )
+    if seqlens_Q is None and seqlens_KV is None:
+        # Not varlen
+        return None, None, 0, 0
+    assert seqlens_Q is not None
+    assert seqlens_KV is not None
+    if not isinstance(seqlens_Q, Tensor) or not isinstance(seqlens_KV, Tensor):
+        raise ValueError("seqlens_Q and seqlens_KV must both be tensors.")
+    if seqlens_Q.device != query.device or seqlens_KV.device != query.device:
+        raise ValueError(
+            "seqlens_Q and seqlens_KV must be on the same device as QKV, but "
+            f"{seqlens_Q.device=}, {seqlens_KV.device=}, {query.device=}."
+        )
+    if seqlens_Q.dtype != torch.int32 or seqlens_KV.dtype != torch.int32:
+        raise ValueError(
+            "seqlens_Q and seqlens_KV must both be torch.int32 tensors, got "
+            f"{seqlens_Q.dtype=}, {seqlens_KV.dtype=}."
+        )
+    if seqlens_Q.dim() != 1 or seqlens_KV.dim() != 1:
+        raise ValueError(
+            "seqlens_Q and seqlens_KV must both be 1-D tensors, got "
+            f"{seqlens_Q.dim()=}, {seqlens_KV.dim()=}."
+        )
+    if seqlens_Q.shape[0] != seqlens_KV.shape[0]:
+        raise ValueError(
+            "seqlens_Q and seqlens_KV must match in size, got "
+            f"{seqlens_Q.shape=}, {seqlens_KV.shape=}."
+        )
+    if seqlens_Q.shape[0] < 1:
+        raise ValueError(
+            "seqlens_Q and seqlens_KV must contain at least one element, got "
+            f"{seqlens_Q.shape=}, {seqlens_KV.shape=}."
+        )
+    if query.shape[0] != 1:
+        raise ValueError(
+            "Variable length attention only supports sequence-packed memory layout "
+            f"(batch = 1), got {query.shape[0]=}."
+        )
+    assert seqlens_Q.dim() == seqlens_KV.dim() == 1
+    assert seqlens_Q.shape[0] == seqlens_KV.shape[0] >= 1
+    assert seqlens_Q.dtype == seqlens_KV.dtype == torch.int32
+    max_seqlen_Q = seqlens_Q.max().item()  # type: ignore
+    max_seqlen_KV = seqlens_KV.max().item()  # type: ignore
+    # NOTE: we have to prepend with 0 manually :(
+    z = torch.tensor([0], dtype=torch.int32, device=seqlens_Q.device)
+    cumulative_seqlen_Q = torch.cat([z, seqlens_Q.cumsum(0).to(torch.int32)], dim=0)
+    cumulative_seqlen_KV = torch.cat([z, seqlens_KV.cumsum(0).to(torch.int32)], dim=0)
+    assert isinstance(max_seqlen_Q, int)
+    assert isinstance(max_seqlen_KV, int)
+    return (
+        cumulative_seqlen_Q,
+        cumulative_seqlen_KV,
+        max_seqlen_Q,
+        max_seqlen_KV,
+    )