diff --git a/build/torch211-cxx11-cu126-x86_64-linux/__init__.py b/build/torch211-cxx11-cu126-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9fe10d400adaa8333f0da9c0c9654864659fbc3b --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/__init__.py @@ -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", +] diff --git a/build/torch211-cxx11-cu126-x86_64-linux/_environment.py b/build/torch211-cxx11-cu126-x86_64-linux/_environment.py new file mode 100644 index 0000000000000000000000000000000000000000..caae85bbe3c5eef992e4c6758e44b382382ebe41 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/_environment.py @@ -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", +] diff --git a/build/torch211-cxx11-cu126-x86_64-linux/_libnatten/__init__.py b/build/torch211-cxx11-cu126-x86_64-linux/_libnatten/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9566f8d4e7d9356415bb2951613dc5ada194dbd9 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/_libnatten/__init__.py @@ -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", +] diff --git a/build/torch211-cxx11-cu126-x86_64-linux/_libnatten/torch_wrappers.py b/build/torch211-cxx11-cu126-x86_64-linux/_libnatten/torch_wrappers.py new file mode 100644 index 0000000000000000000000000000000000000000..16e5b62b2bc5de0810f58d82d70d7fb2e3e56e4a --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/_libnatten/torch_wrappers.py @@ -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", +] diff --git a/build/torch211-cxx11-cu126-x86_64-linux/_natten_cuda_3641131.abi3.so b/build/torch211-cxx11-cu126-x86_64-linux/_natten_cuda_3641131.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..6c138dedcc617ac62ed3ce63e997f52984b8c1b8 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/_natten_cuda_3641131.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a38a8d19b1065ea660148531865732a9ec754be0efb551fcaeef77ce8d71d890 +size 103606392 diff --git a/build/torch211-cxx11-cu126-x86_64-linux/_ops.py b/build/torch211-cxx11-cu126-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..ee2d34306175408457b5e3dca950160c1444f947 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/_ops.py @@ -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}" diff --git a/build/torch211-cxx11-cu126-x86_64-linux/_types.py b/build/torch211-cxx11-cu126-x86_64-linux/_types.py new file mode 100644 index 0000000000000000000000000000000000000000..804980900b74586f5861f473853055c7d9630efd --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/_types.py @@ -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 diff --git a/build/torch211-cxx11-cu126-x86_64-linux/attn_merge.py b/build/torch211-cxx11-cu126-x86_64-linux/attn_merge.py new file mode 100644 index 0000000000000000000000000000000000000000..e71cc799d407df0afa701942ad08b03718ba20bf --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/attn_merge.py @@ -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"] diff --git a/build/torch211-cxx11-cu126-x86_64-linux/backends/__init__.py b/build/torch211-cxx11-cu126-x86_64-linux/backends/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..84b8192a37e048880aba2059a6f47e7d26f6a8ee --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/backends/__init__.py @@ -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", +] diff --git a/build/torch211-cxx11-cu126-x86_64-linux/backends/blackwell_fmha.py b/build/torch211-cxx11-cu126-x86_64-linux/backends/blackwell_fmha.py new file mode 100644 index 0000000000000000000000000000000000000000..9b36e79d8b9242b7c5a20833932721556abd61d5 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/backends/blackwell_fmha.py @@ -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 diff --git a/build/torch211-cxx11-cu126-x86_64-linux/backends/blackwell_fna.py b/build/torch211-cxx11-cu126-x86_64-linux/backends/blackwell_fna.py new file mode 100644 index 0000000000000000000000000000000000000000..5bd6f0ed8dd3da85bd60804e503f2aca27b37e4e --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/backends/blackwell_fna.py @@ -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, + ) diff --git a/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/__init__.py b/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9c2b75f7e5007dd1490e8ff5d4c6413f21a5e39f --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/__init__.py @@ -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) diff --git a/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/checks.py b/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/checks.py new file mode 100644 index 0000000000000000000000000000000000000000..4237d8eba4f2e11f7f7abfc796d53b87801f34a6 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/checks.py @@ -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 diff --git a/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass/__init__.py b/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fdd2d4e18d0d959b16f7c9bd0a8b2f71081845be --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass/__init__.py @@ -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}" + ) diff --git a/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass/backward_knobs.py b/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass/backward_knobs.py new file mode 100644 index 0000000000000000000000000000000000000000..dfb1c10a3f1dfa28597f6eb30e5f925ad4785325 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass/backward_knobs.py @@ -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.") diff --git a/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_backward_128x128.py b/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_backward_128x128.py new file mode 100644 index 0000000000000000000000000000000000000000..af6047ccb59f6923be82d6379b92c39b9bcb7f1a --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_backward_128x128.py @@ -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)), + ], +} diff --git a/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_backward_128x64.py b/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_backward_128x64.py new file mode 100644 index 0000000000000000000000000000000000000000..6a915e9686203919079092eccc743dce88b4107a --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_backward_128x64.py @@ -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)), + ], +} diff --git a/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_backward_64x64.py b/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_backward_64x64.py new file mode 100644 index 0000000000000000000000000000000000000000..8511bd844f34a263d233d4e6a1f6f7c775139e2d --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_backward_64x64.py @@ -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)), + ], +} diff --git a/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_forward_32x128.py b/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_forward_32x128.py new file mode 100644 index 0000000000000000000000000000000000000000..a930aba67df1de878d7e95797bb3b6d2d903e020 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_forward_32x128.py @@ -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)), + ], +} diff --git a/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_forward_64x128.py b/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_forward_64x128.py new file mode 100644 index 0000000000000000000000000000000000000000..1bc1fc78e028a755b755ea8108264eccf0bb9659 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_forward_64x128.py @@ -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)), + ], +} diff --git a/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_forward_64x64.py b/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_forward_64x64.py new file mode 100644 index 0000000000000000000000000000000000000000..c321c5b4e1a96f0cb58894f95dd2826d2d21c950 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_forward_64x64.py @@ -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)), + ], +} diff --git a/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass_blackwell/__init__.py b/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass_blackwell/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4a8c87c6c6d916f58c533fc2a84c66aff8bfe46b --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass_blackwell/__init__.py @@ -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}" + ) diff --git a/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass_hopper/__init__.py b/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass_hopper/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..70b2e3106003ed583050be356b54e32c9b7c61c0 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/cutlass_hopper/__init__.py @@ -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}" + ) diff --git a/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/flex/__init__.py b/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/flex/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c5027701152d3d5976670fd59dc158331c5cadfa --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/backends/configs/flex/__init__.py @@ -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}" + ) diff --git a/build/torch211-cxx11-cu126-x86_64-linux/backends/flex.py b/build/torch211-cxx11-cu126-x86_64-linux/backends/flex.py new file mode 100644 index 0000000000000000000000000000000000000000..0555a7ebdb62e6e4060fbfe3106bed116b44797b --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/backends/flex.py @@ -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, + ) diff --git a/build/torch211-cxx11-cu126-x86_64-linux/backends/fmha.py b/build/torch211-cxx11-cu126-x86_64-linux/backends/fmha.py new file mode 100644 index 0000000000000000000000000000000000000000..0c6851f8bb362a4c7d45db46e61db1ed1dc61d51 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/backends/fmha.py @@ -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 diff --git a/build/torch211-cxx11-cu126-x86_64-linux/backends/fna.py b/build/torch211-cxx11-cu126-x86_64-linux/backends/fna.py new file mode 100644 index 0000000000000000000000000000000000000000..747ecb1f7dcfe032e02e4fa76b8a3a9b70e96bd0 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/backends/fna.py @@ -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, + ) diff --git a/build/torch211-cxx11-cu126-x86_64-linux/backends/hopper_fmha.py b/build/torch211-cxx11-cu126-x86_64-linux/backends/hopper_fmha.py new file mode 100644 index 0000000000000000000000000000000000000000..251b50dc34598111f288ed3ac3d2b414ea23f253 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/backends/hopper_fmha.py @@ -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 diff --git a/build/torch211-cxx11-cu126-x86_64-linux/backends/hopper_fna.py b/build/torch211-cxx11-cu126-x86_64-linux/backends/hopper_fna.py new file mode 100644 index 0000000000000000000000000000000000000000..17bb6d5952675043b66b59fcf92ea9eea54facf9 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/backends/hopper_fna.py @@ -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, + ) diff --git a/build/torch211-cxx11-cu126-x86_64-linux/backends/reference.py b/build/torch211-cxx11-cu126-x86_64-linux/backends/reference.py new file mode 100644 index 0000000000000000000000000000000000000000..d0f34b41a8c9b1a1051360b23d7e20573931bb27 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/backends/reference.py @@ -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, + ) diff --git a/build/torch211-cxx11-cu126-x86_64-linux/context.py b/build/torch211-cxx11-cu126-x86_64-linux/context.py new file mode 100644 index 0000000000000000000000000000000000000000..1793ac984bbe5be383bec2c07b54c770c9969b42 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/context.py @@ -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) diff --git a/build/torch211-cxx11-cu126-x86_64-linux/functional.py b/build/torch211-cxx11-cu126-x86_64-linux/functional.py new file mode 100644 index 0000000000000000000000000000000000000000..944a0c2a7786de250fd5031e933133ea676d2a79 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/functional.py @@ -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, + ) diff --git a/build/torch211-cxx11-cu126-x86_64-linux/metadata.json b/build/torch211-cxx11-cu126-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..a0fa8d3bd8ec9e503fd2cc5f03df62ccf16ee2fd --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/metadata.json @@ -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 + } + } +} \ No newline at end of file diff --git a/build/torch211-cxx11-cu126-x86_64-linux/metadata.json.sigstore b/build/torch211-cxx11-cu126-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..ccd13921b964294af2343dc8d8eb3cb6a031d4f0 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/metadata.json.sigstore @@ -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=="}} \ No newline at end of file diff --git a/build/torch211-cxx11-cu126-x86_64-linux/modules.py b/build/torch211-cxx11-cu126-x86_64-linux/modules.py new file mode 100644 index 0000000000000000000000000000000000000000..222e384f2151c6a32275abcdb5000d5b8f1379d4 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/modules.py @@ -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, + ) diff --git a/build/torch211-cxx11-cu126-x86_64-linux/natten/__init__.py b/build/torch211-cxx11-cu126-x86_64-linux/natten/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/natten/__init__.py @@ -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"))) diff --git a/build/torch211-cxx11-cu126-x86_64-linux/token_permute/__init__.py b/build/torch211-cxx11-cu126-x86_64-linux/token_permute/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..381de37e4e7d4d2e6158668fbcc2ca52f7b388a3 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/token_permute/__init__.py @@ -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", +] diff --git a/build/torch211-cxx11-cu126-x86_64-linux/token_permute/cutlass_impl.py b/build/torch211-cxx11-cu126-x86_64-linux/token_permute/cutlass_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..49ce678f9052679448f843922f16bdec5d8796be --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/token_permute/cutlass_impl.py @@ -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 diff --git a/build/torch211-cxx11-cu126-x86_64-linux/token_permute/frontend.py b/build/torch211-cxx11-cu126-x86_64-linux/token_permute/frontend.py new file mode 100644 index 0000000000000000000000000000000000000000..720aded02e93e31f19a91efbda82cac3db15e4d2 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/token_permute/frontend.py @@ -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 diff --git a/build/torch211-cxx11-cu126-x86_64-linux/token_permute/torch_impl.py b/build/torch211-cxx11-cu126-x86_64-linux/token_permute/torch_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..7f06204e57cb6fc96ccb9e10ac9f798ac2d0e7a4 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/token_permute/torch_impl.py @@ -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", +] diff --git a/build/torch211-cxx11-cu126-x86_64-linux/utils/__init__.py b/build/torch211-cxx11-cu126-x86_64-linux/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..99f6c1e8cca79bfdf04640b8a92602b205a407a0 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/utils/__init__.py @@ -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. +# +################################################################################################# diff --git a/build/torch211-cxx11-cu126-x86_64-linux/utils/checks.py b/build/torch211-cxx11-cu126-x86_64-linux/utils/checks.py new file mode 100644 index 0000000000000000000000000000000000000000..f58db50a581e0bc6b43e8de0fbccd9404beaa1dc --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/utils/checks.py @@ -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, + ) diff --git a/build/torch211-cxx11-cu126-x86_64-linux/utils/device.py b/build/torch211-cxx11-cu126-x86_64-linux/utils/device.py new file mode 100644 index 0000000000000000000000000000000000000000..7997ab474197c78f3d41138f17ae311d6e06b1cf --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/utils/device.py @@ -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 diff --git a/build/torch211-cxx11-cu126-x86_64-linux/utils/dtype.py b/build/torch211-cxx11-cu126-x86_64-linux/utils/dtype.py new file mode 100644 index 0000000000000000000000000000000000000000..c65d9e9e0082c6f14af0b8331216cf26fcb3f28d --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/utils/dtype.py @@ -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] diff --git a/build/torch211-cxx11-cu126-x86_64-linux/utils/environment.py b/build/torch211-cxx11-cu126-x86_64-linux/utils/environment.py new file mode 100644 index 0000000000000000000000000000000000000000..0581aacb1183a5dea228a63859837113429e1a62 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/utils/environment.py @@ -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 diff --git a/build/torch211-cxx11-cu126-x86_64-linux/utils/log.py b/build/torch211-cxx11-cu126-x86_64-linux/utils/log.py new file mode 100644 index 0000000000000000000000000000000000000000..138c3ac8763480c0a381a8af0d913efcfa48b1a1 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/utils/log.py @@ -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) diff --git a/build/torch211-cxx11-cu126-x86_64-linux/utils/tensor.py b/build/torch211-cxx11-cu126-x86_64-linux/utils/tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..32d95acf2eb9a8c6fa29c7d3de3b48b5f14be415 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/utils/tensor.py @@ -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 diff --git a/build/torch211-cxx11-cu126-x86_64-linux/utils/testing.py b/build/torch211-cxx11-cu126-x86_64-linux/utils/testing.py new file mode 100644 index 0000000000000000000000000000000000000000..8ed9de5df24754dc9207a0b39954c8a73d7610a5 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/utils/testing.py @@ -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 diff --git a/build/torch211-cxx11-cu126-x86_64-linux/utils/tuples.py b/build/torch211-cxx11-cu126-x86_64-linux/utils/tuples.py new file mode 100644 index 0000000000000000000000000000000000000000..abf8a862fa4b647aab55fb2982e5dea08609d316 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/utils/tuples.py @@ -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 diff --git a/build/torch211-cxx11-cu126-x86_64-linux/utils/varlen.py b/build/torch211-cxx11-cu126-x86_64-linux/utils/varlen.py new file mode 100644 index 0000000000000000000000000000000000000000..8bc541dbd8e80ddf9923a097ac473b31035174ef --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/utils/varlen.py @@ -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, + ) diff --git a/build/torch211-cxx11-cu126-x86_64-linux/version.py b/build/torch211-cxx11-cu126-x86_64-linux/version.py new file mode 100644 index 0000000000000000000000000000000000000000..fe22ed72b87e6b242b01aae1ffb3de069ae7b0a9 --- /dev/null +++ b/build/torch211-cxx11-cu126-x86_64-linux/version.py @@ -0,0 +1,24 @@ +################################################################################################# +# 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. +# +################################################################################################# + +__version__ = "0.21.7" diff --git a/build/torch211-cxx11-cu128-x86_64-linux/__init__.py b/build/torch211-cxx11-cu128-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9fe10d400adaa8333f0da9c0c9654864659fbc3b --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/__init__.py @@ -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", +] diff --git a/build/torch211-cxx11-cu128-x86_64-linux/_environment.py b/build/torch211-cxx11-cu128-x86_64-linux/_environment.py new file mode 100644 index 0000000000000000000000000000000000000000..caae85bbe3c5eef992e4c6758e44b382382ebe41 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/_environment.py @@ -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", +] diff --git a/build/torch211-cxx11-cu128-x86_64-linux/_libnatten/__init__.py b/build/torch211-cxx11-cu128-x86_64-linux/_libnatten/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9566f8d4e7d9356415bb2951613dc5ada194dbd9 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/_libnatten/__init__.py @@ -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", +] diff --git a/build/torch211-cxx11-cu128-x86_64-linux/_libnatten/torch_wrappers.py b/build/torch211-cxx11-cu128-x86_64-linux/_libnatten/torch_wrappers.py new file mode 100644 index 0000000000000000000000000000000000000000..16e5b62b2bc5de0810f58d82d70d7fb2e3e56e4a --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/_libnatten/torch_wrappers.py @@ -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", +] diff --git a/build/torch211-cxx11-cu128-x86_64-linux/_natten_cuda_3641131.abi3.so b/build/torch211-cxx11-cu128-x86_64-linux/_natten_cuda_3641131.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..6eb62e126f162c1f496f5602ca7e8cfbf868e555 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/_natten_cuda_3641131.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e490ef2cb3a27d7230bbe30a5f64f96ffcbe23cfd56abea3d47ac3d1132ddb3e +size 336844600 diff --git a/build/torch211-cxx11-cu128-x86_64-linux/_ops.py b/build/torch211-cxx11-cu128-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..ee2d34306175408457b5e3dca950160c1444f947 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/_ops.py @@ -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}" diff --git a/build/torch211-cxx11-cu128-x86_64-linux/_types.py b/build/torch211-cxx11-cu128-x86_64-linux/_types.py new file mode 100644 index 0000000000000000000000000000000000000000..804980900b74586f5861f473853055c7d9630efd --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/_types.py @@ -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 diff --git a/build/torch211-cxx11-cu128-x86_64-linux/attn_merge.py b/build/torch211-cxx11-cu128-x86_64-linux/attn_merge.py new file mode 100644 index 0000000000000000000000000000000000000000..e71cc799d407df0afa701942ad08b03718ba20bf --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/attn_merge.py @@ -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"] diff --git a/build/torch211-cxx11-cu128-x86_64-linux/backends/__init__.py b/build/torch211-cxx11-cu128-x86_64-linux/backends/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..84b8192a37e048880aba2059a6f47e7d26f6a8ee --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/backends/__init__.py @@ -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", +] diff --git a/build/torch211-cxx11-cu128-x86_64-linux/backends/blackwell_fmha.py b/build/torch211-cxx11-cu128-x86_64-linux/backends/blackwell_fmha.py new file mode 100644 index 0000000000000000000000000000000000000000..9b36e79d8b9242b7c5a20833932721556abd61d5 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/backends/blackwell_fmha.py @@ -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 diff --git a/build/torch211-cxx11-cu128-x86_64-linux/backends/blackwell_fna.py b/build/torch211-cxx11-cu128-x86_64-linux/backends/blackwell_fna.py new file mode 100644 index 0000000000000000000000000000000000000000..5bd6f0ed8dd3da85bd60804e503f2aca27b37e4e --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/backends/blackwell_fna.py @@ -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, + ) diff --git a/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/__init__.py b/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9c2b75f7e5007dd1490e8ff5d4c6413f21a5e39f --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/__init__.py @@ -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) diff --git a/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/checks.py b/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/checks.py new file mode 100644 index 0000000000000000000000000000000000000000..4237d8eba4f2e11f7f7abfc796d53b87801f34a6 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/checks.py @@ -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 diff --git a/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass/__init__.py b/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fdd2d4e18d0d959b16f7c9bd0a8b2f71081845be --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass/__init__.py @@ -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}" + ) diff --git a/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass/backward_knobs.py b/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass/backward_knobs.py new file mode 100644 index 0000000000000000000000000000000000000000..dfb1c10a3f1dfa28597f6eb30e5f925ad4785325 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass/backward_knobs.py @@ -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.") diff --git a/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass/fna_backward_128x128.py b/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass/fna_backward_128x128.py new file mode 100644 index 0000000000000000000000000000000000000000..af6047ccb59f6923be82d6379b92c39b9bcb7f1a --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass/fna_backward_128x128.py @@ -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)), + ], +} diff --git a/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass/fna_backward_128x64.py b/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass/fna_backward_128x64.py new file mode 100644 index 0000000000000000000000000000000000000000..6a915e9686203919079092eccc743dce88b4107a --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass/fna_backward_128x64.py @@ -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)), + ], +} diff --git a/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass/fna_backward_64x64.py b/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass/fna_backward_64x64.py new file mode 100644 index 0000000000000000000000000000000000000000..8511bd844f34a263d233d4e6a1f6f7c775139e2d --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass/fna_backward_64x64.py @@ -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)), + ], +} diff --git a/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass/fna_forward_32x128.py b/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass/fna_forward_32x128.py new file mode 100644 index 0000000000000000000000000000000000000000..a930aba67df1de878d7e95797bb3b6d2d903e020 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass/fna_forward_32x128.py @@ -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)), + ], +} diff --git a/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass/fna_forward_64x128.py b/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass/fna_forward_64x128.py new file mode 100644 index 0000000000000000000000000000000000000000..1bc1fc78e028a755b755ea8108264eccf0bb9659 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass/fna_forward_64x128.py @@ -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)), + ], +} diff --git a/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass/fna_forward_64x64.py b/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass/fna_forward_64x64.py new file mode 100644 index 0000000000000000000000000000000000000000..c321c5b4e1a96f0cb58894f95dd2826d2d21c950 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass/fna_forward_64x64.py @@ -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)), + ], +} diff --git a/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass_blackwell/__init__.py b/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass_blackwell/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4a8c87c6c6d916f58c533fc2a84c66aff8bfe46b --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass_blackwell/__init__.py @@ -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}" + ) diff --git a/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass_hopper/__init__.py b/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass_hopper/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..70b2e3106003ed583050be356b54e32c9b7c61c0 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/cutlass_hopper/__init__.py @@ -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}" + ) diff --git a/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/flex/__init__.py b/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/flex/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c5027701152d3d5976670fd59dc158331c5cadfa --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/backends/configs/flex/__init__.py @@ -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}" + ) diff --git a/build/torch211-cxx11-cu128-x86_64-linux/backends/flex.py b/build/torch211-cxx11-cu128-x86_64-linux/backends/flex.py new file mode 100644 index 0000000000000000000000000000000000000000..0555a7ebdb62e6e4060fbfe3106bed116b44797b --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/backends/flex.py @@ -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, + ) diff --git a/build/torch211-cxx11-cu128-x86_64-linux/backends/fmha.py b/build/torch211-cxx11-cu128-x86_64-linux/backends/fmha.py new file mode 100644 index 0000000000000000000000000000000000000000..0c6851f8bb362a4c7d45db46e61db1ed1dc61d51 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/backends/fmha.py @@ -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 diff --git a/build/torch211-cxx11-cu128-x86_64-linux/backends/fna.py b/build/torch211-cxx11-cu128-x86_64-linux/backends/fna.py new file mode 100644 index 0000000000000000000000000000000000000000..747ecb1f7dcfe032e02e4fa76b8a3a9b70e96bd0 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/backends/fna.py @@ -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, + ) diff --git a/build/torch211-cxx11-cu128-x86_64-linux/backends/hopper_fmha.py b/build/torch211-cxx11-cu128-x86_64-linux/backends/hopper_fmha.py new file mode 100644 index 0000000000000000000000000000000000000000..251b50dc34598111f288ed3ac3d2b414ea23f253 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/backends/hopper_fmha.py @@ -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 diff --git a/build/torch211-cxx11-cu128-x86_64-linux/backends/hopper_fna.py b/build/torch211-cxx11-cu128-x86_64-linux/backends/hopper_fna.py new file mode 100644 index 0000000000000000000000000000000000000000..17bb6d5952675043b66b59fcf92ea9eea54facf9 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/backends/hopper_fna.py @@ -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, + ) diff --git a/build/torch211-cxx11-cu128-x86_64-linux/backends/reference.py b/build/torch211-cxx11-cu128-x86_64-linux/backends/reference.py new file mode 100644 index 0000000000000000000000000000000000000000..d0f34b41a8c9b1a1051360b23d7e20573931bb27 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/backends/reference.py @@ -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, + ) diff --git a/build/torch211-cxx11-cu128-x86_64-linux/context.py b/build/torch211-cxx11-cu128-x86_64-linux/context.py new file mode 100644 index 0000000000000000000000000000000000000000..1793ac984bbe5be383bec2c07b54c770c9969b42 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/context.py @@ -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) diff --git a/build/torch211-cxx11-cu128-x86_64-linux/functional.py b/build/torch211-cxx11-cu128-x86_64-linux/functional.py new file mode 100644 index 0000000000000000000000000000000000000000..944a0c2a7786de250fd5031e933133ea676d2a79 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/functional.py @@ -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, + ) diff --git a/build/torch211-cxx11-cu128-x86_64-linux/metadata.json b/build/torch211-cxx11-cu128-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..b68b3dc43c648af8d5122e709758be60374a9d7b --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/metadata.json @@ -0,0 +1,84 @@ +{ + "name": "natten", + "id": "_natten_cuda_3641131", + "version": 1, + "license": "MIT", + "upstream": "https://github.com/SHI-Labs/NATTEN", + "python-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "10.0", + "10.0a", + "12.0", + "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": "5JDvLLOifXIwu+MKX2T5b/y+I8/Var6j1HrD0RMt2z4=", + "_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 + } + } +} \ No newline at end of file diff --git a/build/torch211-cxx11-cu128-x86_64-linux/metadata.json.sigstore b/build/torch211-cxx11-cu128-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..4437fec3042bc3563091c0d072e9e699fb672019 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtGgAwIBAgIUedd/FFlf012F73eWOvCOXcO2vFIwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNzI5MDkyMTM2WhcNMjYwNzI5MDkzMTM2WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEA5mIKgyrGeIM+OlJz6lL1vnd9Ft4H15DUpw+edwaKoLeuWzs9sSICYDhY+fd/OXdoFIf7BF7NpnoG5ut+Y3ESaOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUdw+F2fgIOD+qmLMnzt4TtNdEwFgwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoMjBhNjEzYzEwZDliMWM5Y2I0M2E1MWU4NTBlOWZhMDQ2OWE5NzE5OTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoMjBhNjEzYzEwZDliMWM5Y2I0M2E1MWU4NTBlOWZhMDQ2OWE5NzE5OTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoMjBhNjEzYzEwZDliMWM5Y2I0M2E1MWU4NTBlOWZhMDQ2OWE5NzE5OTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDIwYTYxM2MxMGQ5YjFjOWNiNDNhNTFlODUwZTlmYTA0NjlhOTcxOTkwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzA0MTgyNTI4NjkvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABn60t71UAAAQDAEcwRQIhAKYd7D6SDB0boYLstrnptUiHC6phxT/9Y+l4FHz/uzgYAiApEiy4A75xeVqUEVoFb1rwpuuRxx+Q3XUXivDnBbMuHDAKBggqhkjOPQQDAwNoADBlAjAA+2miTkJhChNbvrWtJBth1s6KbxeuSdHaKTfLiKfiFlfdwWOXlg/y5cOBI79dBk8CMQDpDqihGm+QBBbOR8u9NKSa3G2Q+D8z4CsdvTnNbPulMMhBTF1e0ViCd2BwhO2yjW8="}, "tlogEntries":[{"logIndex":"2280149266", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1785316896", "inclusionPromise":{"signedEntryTimestamp":"MEUCIQCLjrULCQNUxqnVldrVPRTEf5NzrWSGtpGCmjvdKKRj8AIgIg38h3MBElw79VgKQRbEJePNJ4szz7tIcZfvyJpRwUE="}, "inclusionProof":{"logIndex":"2158245004", "rootHash":"G7otfalMZT0tUyySbpGp/j+9LIwIfPhbqGG4Cxep7cQ=", "treeSize":"2158245010", "hashes":["C6g/+HUEJzh4t0p7zTLJV7dMaimmxfZYmyaDLiL8wmE=", "bTkvleWVGhwLNC8oA4tqrma07lxQ2rNEp+F3au2u9zc=", "JuFKkP048H9820X1HNoLpb7Iu9v6Ke+iYzLatlV/pVw=", "aSIN7kR6a/eLoHwT7eRsyVAOhV1SQIcBWXZvo0r6YpM=", "voOpzARDbsLIxR2qAHOIHYLwFkywk98DcRDRMM9x/10=", "/5uu/jw8GTRTqEZ0auUV9lv+Zg0twg6hHe7ZChWeiJk=", "DwNueYXiyuVjob7c8CRNRN57ioovkJTOeOeOQIi6/pI=", "bXYdybsl/2BK7Yled3FQJEE3Hk0TeGcWnXUxOkVZv8w=", "xrwx0yHkH5ZmCgeWs/jLZZ8RcejqpeuUeAWNJkAYrHk=", "8MN6j2GHUjPnNtIwpas9l/XjdwE0/KPUhFDlrVkngko=", "b+xUZfuENQxvSOJxzNvYvRG8eVphfszPpZmuf4/cQ6c=", "OVsvZCKnWA+498QUIaQCtitUT6huDbC7SmhH1l8MxXI=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2158245010\nG7otfalMZT0tUyySbpGp/j+9LIwIfPhbqGG4Cxep7cQ=\n\n— rekor.sigstore.dev wNI9ajBFAiA8NS6Xod2pYjCQ98Px40MverUVjTzz7/RZ9fKYCDim0QIhANiogbH/lZMLmdia5M9uFp05aRsYVDs6KkH5bMFoK7zn\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiIyM2Q1ZWVjMjY2M2FhYTRmNjM5OGQxNWVkNWM3NDdhNTFiZjNhNjliNGY5Yzc0MmQ3OTM1MTczOGEwNTA0Y2I5In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FWUNJUUNOUFBFMXc1OHJNbzdGdklMbFI1emtZbkp3QlJXMEpjMkdKUCtDYlBVSFdRSWhBSlRjS1lJbnVQMU1pVjdleVZQazMzczI0WFJRZDAyaWVteDNtZUFxV3dTMiIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5SSFowRjNTVUpCWjBsVlpXUmtMMFpHYkdZd01USkdOek5sVjA5MlEwOVlZMDh5ZGtaSmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDU2U1RWTlJHdDVUVlJOTWxkb1kwNU5hbGwzVG5wSk5VMUVhM3BOVkUweVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZCTlcxSlMyZDVja2RsU1UwclQyeEtlalpzVERGMmJtUTVSblEwU0RFMVJGVndkeXNLWldSM1lVdHZUR1YxVjNwek9YTlRTVU5aUkdoWksyWmtMMDlZWkc5R1NXWTNRa1kzVG5CdWIwYzFkWFFyV1RORlUyRlBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZrZHl0R0NqSm1aMGxQUkN0eGJVeE5ibnAwTkZSMFRtUkZkMFpuZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU5ha0pvVG1wRmVsbDZSWGRhUkd4cFRWZE5OVmt5U1RCTk1rVXhDazFYVlRST1ZFSnNUMWRhYUUxRVVUSlBWMFUxVG5wRk5VOVVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDFxUW1oT2FrVjZXWHBGZDFwRWJHbE5WMDAxV1RKSk1FMHlSVEZOVjFVMFRsUkNiRTlYV21nS1RVUlJNazlYUlRWT2VrVTFUMVJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDFxUW1oT2FrVjZXWHBGZDFwRWJHbE5WMDAxV1RKSk1FMHlSVEVLVFZkVk5FNVVRbXhQVjFwb1RVUlJNazlYUlRWT2VrVTFUMVJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUkpkMWxVV1hnS1RUSk5lRTFIVVRWWmFrWnFUMWRPYVU1RVRtaE9WRVpzVDBSVmQxcFViRzFaVkVFd1RtcHNhRTlVWTNoUFZHdDNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWtFd1RWUm5lVTVVU1RST2FtdDJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNDJNSFEzTVZWQlFVRlJSQXBCUldOM1VsRkphRUZMV1dRM1JEWlRSRUl3WW05WlRITjBjbTV3ZEZWcFNFTTJjR2g0VkM4NVdTdHNORVpJZWk5MWVtZFpRV2xCY0VWcGVUUkJOelY0Q21WV2NWVkZWbTlHWWpGeWQzQjFkVko0ZUN0Uk0xaFZXR2wyUkc1Q1lrMTFTRVJCUzBKblozRm9hMnBQVUZGUlJFRjNUbTlCUkVKc1FXcEJRU3N5YldrS1ZHdEthRU5vVG1KMmNsZDBTa0owYURGek5rdGllR1YxVTJSSVlVdFVaa3hwUzJacFJteG1aSGRYVDFoc1p5OTVOV05QUWtrM09XUkNhemhEVFZGRWNBcEVjV2xvUjIwclVVSkNZazlTT0hVNVRrdFRZVE5ITWxFclJEaDZORU56WkhaVWJrNWlVSFZzVFUxb1FsUkdNV1V3Vm1sRFpESkNkMmhQTW5scVZ6ZzlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyDADAgEAMIICvwYJKoZIhvcNAQcCoIICsDCCAqwCAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgv5aDz6NF/Qj4MdV9Q7/T/YNekTgZl5fT6RAC/I/tHyQCFA7JD8g3J+pD7RrP0QKb91cnoZULGA8yMDI2MDcyOTA5MjEzNlowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdowggHWAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwNzI5MDkyMTM2WjAvBgkqhkiG9w0BCQQxIgQguddgn/zH4nmhS7isUxxnpw+2H7734RK3VJ24axgYfnAwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGYwZAIwHbmW+tbtvVTE4a9NMSPoJzkm8AUthWRfhr+xL4fy6k/YpTEvXJEM4BEbd5fuqLdOAjAbtZ54KTbqshFdhga6J0An95VHlvLbaQ76/JP+CIbjn4psMu+65XwwjyVC2wLs14Y="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"I9XuwmY6qk9jmNFe1cdHpRvzpptPnHQteTUXOKBQTLk="}, "signature":"MEYCIQCNPPE1w58rMo7FvILlR5zkYnJwBRW0Jc2GJP+CbPUHWQIhAJTcKYInuP1MiV7eyVPk33s24XRQd02iemx3meAqWwS2"}} \ No newline at end of file diff --git a/build/torch211-cxx11-cu128-x86_64-linux/modules.py b/build/torch211-cxx11-cu128-x86_64-linux/modules.py new file mode 100644 index 0000000000000000000000000000000000000000..222e384f2151c6a32275abcdb5000d5b8f1379d4 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/modules.py @@ -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, + ) diff --git a/build/torch211-cxx11-cu128-x86_64-linux/natten/__init__.py b/build/torch211-cxx11-cu128-x86_64-linux/natten/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/natten/__init__.py @@ -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"))) diff --git a/build/torch211-cxx11-cu128-x86_64-linux/token_permute/__init__.py b/build/torch211-cxx11-cu128-x86_64-linux/token_permute/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..381de37e4e7d4d2e6158668fbcc2ca52f7b388a3 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/token_permute/__init__.py @@ -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", +] diff --git a/build/torch211-cxx11-cu128-x86_64-linux/token_permute/cutlass_impl.py b/build/torch211-cxx11-cu128-x86_64-linux/token_permute/cutlass_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..49ce678f9052679448f843922f16bdec5d8796be --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/token_permute/cutlass_impl.py @@ -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 diff --git a/build/torch211-cxx11-cu128-x86_64-linux/token_permute/frontend.py b/build/torch211-cxx11-cu128-x86_64-linux/token_permute/frontend.py new file mode 100644 index 0000000000000000000000000000000000000000..720aded02e93e31f19a91efbda82cac3db15e4d2 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/token_permute/frontend.py @@ -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 diff --git a/build/torch211-cxx11-cu128-x86_64-linux/token_permute/torch_impl.py b/build/torch211-cxx11-cu128-x86_64-linux/token_permute/torch_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..7f06204e57cb6fc96ccb9e10ac9f798ac2d0e7a4 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/token_permute/torch_impl.py @@ -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", +] diff --git a/build/torch211-cxx11-cu128-x86_64-linux/utils/__init__.py b/build/torch211-cxx11-cu128-x86_64-linux/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..99f6c1e8cca79bfdf04640b8a92602b205a407a0 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/utils/__init__.py @@ -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. +# +################################################################################################# diff --git a/build/torch211-cxx11-cu128-x86_64-linux/utils/checks.py b/build/torch211-cxx11-cu128-x86_64-linux/utils/checks.py new file mode 100644 index 0000000000000000000000000000000000000000..f58db50a581e0bc6b43e8de0fbccd9404beaa1dc --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/utils/checks.py @@ -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, + ) diff --git a/build/torch211-cxx11-cu128-x86_64-linux/utils/device.py b/build/torch211-cxx11-cu128-x86_64-linux/utils/device.py new file mode 100644 index 0000000000000000000000000000000000000000..7997ab474197c78f3d41138f17ae311d6e06b1cf --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/utils/device.py @@ -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 diff --git a/build/torch211-cxx11-cu128-x86_64-linux/utils/dtype.py b/build/torch211-cxx11-cu128-x86_64-linux/utils/dtype.py new file mode 100644 index 0000000000000000000000000000000000000000..c65d9e9e0082c6f14af0b8331216cf26fcb3f28d --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/utils/dtype.py @@ -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] diff --git a/build/torch211-cxx11-cu128-x86_64-linux/utils/environment.py b/build/torch211-cxx11-cu128-x86_64-linux/utils/environment.py new file mode 100644 index 0000000000000000000000000000000000000000..0581aacb1183a5dea228a63859837113429e1a62 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/utils/environment.py @@ -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 diff --git a/build/torch211-cxx11-cu128-x86_64-linux/utils/log.py b/build/torch211-cxx11-cu128-x86_64-linux/utils/log.py new file mode 100644 index 0000000000000000000000000000000000000000..138c3ac8763480c0a381a8af0d913efcfa48b1a1 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/utils/log.py @@ -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) diff --git a/build/torch211-cxx11-cu128-x86_64-linux/utils/tensor.py b/build/torch211-cxx11-cu128-x86_64-linux/utils/tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..32d95acf2eb9a8c6fa29c7d3de3b48b5f14be415 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/utils/tensor.py @@ -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 diff --git a/build/torch211-cxx11-cu128-x86_64-linux/utils/testing.py b/build/torch211-cxx11-cu128-x86_64-linux/utils/testing.py new file mode 100644 index 0000000000000000000000000000000000000000..8ed9de5df24754dc9207a0b39954c8a73d7610a5 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/utils/testing.py @@ -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 diff --git a/build/torch211-cxx11-cu128-x86_64-linux/utils/tuples.py b/build/torch211-cxx11-cu128-x86_64-linux/utils/tuples.py new file mode 100644 index 0000000000000000000000000000000000000000..abf8a862fa4b647aab55fb2982e5dea08609d316 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/utils/tuples.py @@ -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 diff --git a/build/torch211-cxx11-cu128-x86_64-linux/utils/varlen.py b/build/torch211-cxx11-cu128-x86_64-linux/utils/varlen.py new file mode 100644 index 0000000000000000000000000000000000000000..8bc541dbd8e80ddf9923a097ac473b31035174ef --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/utils/varlen.py @@ -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, + ) diff --git a/build/torch211-cxx11-cu128-x86_64-linux/version.py b/build/torch211-cxx11-cu128-x86_64-linux/version.py new file mode 100644 index 0000000000000000000000000000000000000000..fe22ed72b87e6b242b01aae1ffb3de069ae7b0a9 --- /dev/null +++ b/build/torch211-cxx11-cu128-x86_64-linux/version.py @@ -0,0 +1,24 @@ +################################################################################################# +# 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. +# +################################################################################################# + +__version__ = "0.21.7" diff --git a/build/torch211-cxx11-cu130-x86_64-linux/__init__.py b/build/torch211-cxx11-cu130-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9fe10d400adaa8333f0da9c0c9654864659fbc3b --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/__init__.py @@ -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", +] diff --git a/build/torch211-cxx11-cu130-x86_64-linux/_environment.py b/build/torch211-cxx11-cu130-x86_64-linux/_environment.py new file mode 100644 index 0000000000000000000000000000000000000000..caae85bbe3c5eef992e4c6758e44b382382ebe41 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/_environment.py @@ -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", +] diff --git a/build/torch211-cxx11-cu130-x86_64-linux/_libnatten/__init__.py b/build/torch211-cxx11-cu130-x86_64-linux/_libnatten/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9566f8d4e7d9356415bb2951613dc5ada194dbd9 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/_libnatten/__init__.py @@ -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", +] diff --git a/build/torch211-cxx11-cu130-x86_64-linux/_libnatten/torch_wrappers.py b/build/torch211-cxx11-cu130-x86_64-linux/_libnatten/torch_wrappers.py new file mode 100644 index 0000000000000000000000000000000000000000..16e5b62b2bc5de0810f58d82d70d7fb2e3e56e4a --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/_libnatten/torch_wrappers.py @@ -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", +] diff --git a/build/torch211-cxx11-cu130-x86_64-linux/_natten_cuda_3641131.abi3.so b/build/torch211-cxx11-cu130-x86_64-linux/_natten_cuda_3641131.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..0937ca4184b9e599bc243bbdc1f84b5f9c4e9323 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/_natten_cuda_3641131.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0795355dcd1d43462f6b54481879723100288b0c20a2c09cd4bc7d2d46687139 +size 164385840 diff --git a/build/torch211-cxx11-cu130-x86_64-linux/_ops.py b/build/torch211-cxx11-cu130-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..ee2d34306175408457b5e3dca950160c1444f947 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/_ops.py @@ -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}" diff --git a/build/torch211-cxx11-cu130-x86_64-linux/_types.py b/build/torch211-cxx11-cu130-x86_64-linux/_types.py new file mode 100644 index 0000000000000000000000000000000000000000..804980900b74586f5861f473853055c7d9630efd --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/_types.py @@ -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 diff --git a/build/torch211-cxx11-cu130-x86_64-linux/attn_merge.py b/build/torch211-cxx11-cu130-x86_64-linux/attn_merge.py new file mode 100644 index 0000000000000000000000000000000000000000..e71cc799d407df0afa701942ad08b03718ba20bf --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/attn_merge.py @@ -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"] diff --git a/build/torch211-cxx11-cu130-x86_64-linux/backends/__init__.py b/build/torch211-cxx11-cu130-x86_64-linux/backends/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..84b8192a37e048880aba2059a6f47e7d26f6a8ee --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/backends/__init__.py @@ -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", +] diff --git a/build/torch211-cxx11-cu130-x86_64-linux/backends/blackwell_fmha.py b/build/torch211-cxx11-cu130-x86_64-linux/backends/blackwell_fmha.py new file mode 100644 index 0000000000000000000000000000000000000000..9b36e79d8b9242b7c5a20833932721556abd61d5 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/backends/blackwell_fmha.py @@ -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 diff --git a/build/torch211-cxx11-cu130-x86_64-linux/backends/blackwell_fna.py b/build/torch211-cxx11-cu130-x86_64-linux/backends/blackwell_fna.py new file mode 100644 index 0000000000000000000000000000000000000000..5bd6f0ed8dd3da85bd60804e503f2aca27b37e4e --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/backends/blackwell_fna.py @@ -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, + ) diff --git a/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/__init__.py b/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9c2b75f7e5007dd1490e8ff5d4c6413f21a5e39f --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/__init__.py @@ -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) diff --git a/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/checks.py b/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/checks.py new file mode 100644 index 0000000000000000000000000000000000000000..4237d8eba4f2e11f7f7abfc796d53b87801f34a6 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/checks.py @@ -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 diff --git a/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass/__init__.py b/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fdd2d4e18d0d959b16f7c9bd0a8b2f71081845be --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass/__init__.py @@ -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}" + ) diff --git a/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass/backward_knobs.py b/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass/backward_knobs.py new file mode 100644 index 0000000000000000000000000000000000000000..dfb1c10a3f1dfa28597f6eb30e5f925ad4785325 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass/backward_knobs.py @@ -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.") diff --git a/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_backward_128x128.py b/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_backward_128x128.py new file mode 100644 index 0000000000000000000000000000000000000000..af6047ccb59f6923be82d6379b92c39b9bcb7f1a --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_backward_128x128.py @@ -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)), + ], +} diff --git a/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_backward_128x64.py b/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_backward_128x64.py new file mode 100644 index 0000000000000000000000000000000000000000..6a915e9686203919079092eccc743dce88b4107a --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_backward_128x64.py @@ -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)), + ], +} diff --git a/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_backward_64x64.py b/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_backward_64x64.py new file mode 100644 index 0000000000000000000000000000000000000000..8511bd844f34a263d233d4e6a1f6f7c775139e2d --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_backward_64x64.py @@ -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)), + ], +} diff --git a/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_forward_32x128.py b/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_forward_32x128.py new file mode 100644 index 0000000000000000000000000000000000000000..a930aba67df1de878d7e95797bb3b6d2d903e020 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_forward_32x128.py @@ -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)), + ], +} diff --git a/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_forward_64x128.py b/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_forward_64x128.py new file mode 100644 index 0000000000000000000000000000000000000000..1bc1fc78e028a755b755ea8108264eccf0bb9659 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_forward_64x128.py @@ -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)), + ], +} diff --git a/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_forward_64x64.py b/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_forward_64x64.py new file mode 100644 index 0000000000000000000000000000000000000000..c321c5b4e1a96f0cb58894f95dd2826d2d21c950 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_forward_64x64.py @@ -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)), + ], +} diff --git a/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass_blackwell/__init__.py b/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass_blackwell/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4a8c87c6c6d916f58c533fc2a84c66aff8bfe46b --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass_blackwell/__init__.py @@ -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}" + ) diff --git a/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass_hopper/__init__.py b/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass_hopper/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..70b2e3106003ed583050be356b54e32c9b7c61c0 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/cutlass_hopper/__init__.py @@ -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}" + ) diff --git a/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/flex/__init__.py b/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/flex/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c5027701152d3d5976670fd59dc158331c5cadfa --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/backends/configs/flex/__init__.py @@ -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}" + ) diff --git a/build/torch211-cxx11-cu130-x86_64-linux/backends/flex.py b/build/torch211-cxx11-cu130-x86_64-linux/backends/flex.py new file mode 100644 index 0000000000000000000000000000000000000000..0555a7ebdb62e6e4060fbfe3106bed116b44797b --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/backends/flex.py @@ -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, + ) diff --git a/build/torch211-cxx11-cu130-x86_64-linux/backends/fmha.py b/build/torch211-cxx11-cu130-x86_64-linux/backends/fmha.py new file mode 100644 index 0000000000000000000000000000000000000000..0c6851f8bb362a4c7d45db46e61db1ed1dc61d51 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/backends/fmha.py @@ -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 diff --git a/build/torch211-cxx11-cu130-x86_64-linux/backends/fna.py b/build/torch211-cxx11-cu130-x86_64-linux/backends/fna.py new file mode 100644 index 0000000000000000000000000000000000000000..747ecb1f7dcfe032e02e4fa76b8a3a9b70e96bd0 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/backends/fna.py @@ -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, + ) diff --git a/build/torch211-cxx11-cu130-x86_64-linux/backends/hopper_fmha.py b/build/torch211-cxx11-cu130-x86_64-linux/backends/hopper_fmha.py new file mode 100644 index 0000000000000000000000000000000000000000..251b50dc34598111f288ed3ac3d2b414ea23f253 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/backends/hopper_fmha.py @@ -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 diff --git a/build/torch211-cxx11-cu130-x86_64-linux/backends/hopper_fna.py b/build/torch211-cxx11-cu130-x86_64-linux/backends/hopper_fna.py new file mode 100644 index 0000000000000000000000000000000000000000..17bb6d5952675043b66b59fcf92ea9eea54facf9 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/backends/hopper_fna.py @@ -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, + ) diff --git a/build/torch211-cxx11-cu130-x86_64-linux/backends/reference.py b/build/torch211-cxx11-cu130-x86_64-linux/backends/reference.py new file mode 100644 index 0000000000000000000000000000000000000000..d0f34b41a8c9b1a1051360b23d7e20573931bb27 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/backends/reference.py @@ -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, + ) diff --git a/build/torch211-cxx11-cu130-x86_64-linux/context.py b/build/torch211-cxx11-cu130-x86_64-linux/context.py new file mode 100644 index 0000000000000000000000000000000000000000..1793ac984bbe5be383bec2c07b54c770c9969b42 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/context.py @@ -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) diff --git a/build/torch211-cxx11-cu130-x86_64-linux/functional.py b/build/torch211-cxx11-cu130-x86_64-linux/functional.py new file mode 100644 index 0000000000000000000000000000000000000000..944a0c2a7786de250fd5031e933133ea676d2a79 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/functional.py @@ -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, + ) diff --git a/build/torch211-cxx11-cu130-x86_64-linux/metadata.json b/build/torch211-cxx11-cu130-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..6f09bb97cdbb55fe11c463588026e3c92a3759cf --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/metadata.json @@ -0,0 +1,84 @@ +{ + "name": "natten", + "id": "_natten_cuda_3641131", + "version": 1, + "license": "MIT", + "upstream": "https://github.com/SHI-Labs/NATTEN", + "python-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "10.0", + "10.0a", + "12.0", + "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": "B5U1Xc0dQ0Yva1RIGHlyMQAoiwwgosCc1Lx9LUZocTk=", + "_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 + } + } +} \ No newline at end of file diff --git a/build/torch211-cxx11-cu130-x86_64-linux/metadata.json.sigstore b/build/torch211-cxx11-cu130-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..2b40325de1c737a8e6e83191a17326384a78017f --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtGgAwIBAgIUSgjAiPhtOZ2+o75crTAYWQ3ROzcwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNzI5MDkyMTMxWhcNMjYwNzI5MDkzMTMxWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEFq0Fxyhh8CRvolRqYiWyCNqp1FQwNJ/HPQE01VyizqIbm+PLx8PFG1NoWSdMYCfHc2ldOfkaNEQ52shjz789vKOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUJ0RgWk5i4X73+eBpg9io4nxVvWUwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoMjBhNjEzYzEwZDliMWM5Y2I0M2E1MWU4NTBlOWZhMDQ2OWE5NzE5OTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoMjBhNjEzYzEwZDliMWM5Y2I0M2E1MWU4NTBlOWZhMDQ2OWE5NzE5OTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoMjBhNjEzYzEwZDliMWM5Y2I0M2E1MWU4NTBlOWZhMDQ2OWE5NzE5OTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDIwYTYxM2MxMGQ5YjFjOWNiNDNhNTFlODUwZTlmYTA0NjlhOTcxOTkwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzA0MTgyNTI4NjkvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABn60t3K4AAAQDAEcwRQIhANLGCl11bIJEeqVTKmSXErKRT6S6Pd3kxGhxA3jo72PHAiAeO4ovjo7VgnYLTs3V+iW7xGL8xrSwtmNLmx7zqtygpzAKBggqhkjOPQQDAwNoADBlAjEArzM/D36kb/lxKa+KcfuUDa9NtbZpS3ObGDJ63edXN+zQXnVNpLark/G4Yxarqgt5AjBiN237Vd5YtKxCG+5WdwL3Mr6Aq8URo2D4GqrUC+Y8W/o+BIUGcBPG7LqMcnKaZvo="}, "tlogEntries":[{"logIndex":"2280149186", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1785316892", "inclusionPromise":{"signedEntryTimestamp":"MEQCIHV/rNd8/LD7G+ywjVXkdqfTdsuBg4crPuCxxR7wM+KCAiBlkHEJXTJdHr3H0aJAqB2gMWK+XQBsEWWHdY23k+OimQ=="}, "inclusionProof":{"logIndex":"2158244924", "rootHash":"kF7x+Io6pptvdcW2eES3bdHS34lTMvDq0zNGVM8c3BE=", "treeSize":"2158244929", "hashes":["B6qMSTynRVpXOSt9VAeGQDfOu7/KrAtDFVYTlshuc0Q=", "3zYNADgEVNJfJdMz8TNNuQfr9D9rntq/yx9+6rHoVDw=", "0tTq6Kurk9dfVTtd1tqG5jK64V/F328Zsw6rf4a7UHs=", "Yu0V+0ctftZD19HNwJBnZtgxoHSY0vLI/m2Bqz8VNJo=", "UJRaxDL1bVMOP7NdSwdbbd5d0R48HihLn9dsmYERc/Y=", "3OO6p2Eqxin9xrn0cYAX/MZ9P+PwDc1v9Ehn6Hsltg0=", "m5w60olqYCcXxnx27ADHz/VqnUq3wmzyrcwAUTsk9SI=", "DwNueYXiyuVjob7c8CRNRN57ioovkJTOeOeOQIi6/pI=", "bXYdybsl/2BK7Yled3FQJEE3Hk0TeGcWnXUxOkVZv8w=", "xrwx0yHkH5ZmCgeWs/jLZZ8RcejqpeuUeAWNJkAYrHk=", "8MN6j2GHUjPnNtIwpas9l/XjdwE0/KPUhFDlrVkngko=", "b+xUZfuENQxvSOJxzNvYvRG8eVphfszPpZmuf4/cQ6c=", "OVsvZCKnWA+498QUIaQCtitUT6huDbC7SmhH1l8MxXI=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2158244929\nkF7x+Io6pptvdcW2eES3bdHS34lTMvDq0zNGVM8c3BE=\n\n— rekor.sigstore.dev wNI9ajBFAiEA/jzqZDOo5Nvo5/AWccl6QU2AKHwoCEKFWvIuFPN8lb4CIGZfHat2YhKQ4ERJVrgDsv+lLKQ4OANqT11fNScBhTuX\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiIxZGM0NjM1MGYyZGJkY2FiODcxMWQwMzZhOGIzMGRhYzUwYzgzMzFkZGQ2Y2IxNmNkMGI5YTgxODY1MTA1MTdhIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUNZSWV2cHFjTkpoUzNZQVMrZ3J5eTlGd2lrZ3g0N2h1MHdYUVlUOGIrMmd3SWdJazJzSk1Ld1VxakFaUEpuTE1Nc3I4Y015THh4dTEyeFNVRmFYaW1GSjkwPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5SSFowRjNTVUpCWjBsVlUyZHFRV2xRYUhSUFdqSXJiemMxWTNKVVFWbFhVVE5TVDNwamQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDU2U1RWTlJHdDVUVlJOZUZkb1kwNU5hbGwzVG5wSk5VMUVhM3BOVkUxNFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZHY1RCR2VIbG9hRGhEVW5admJGSnhXV2xYZVVOT2NYQXhSbEYzVGtvdlNGQlJSVEFLTVZaNWFYcHhTV0p0SzFCTWVEaFFSa2N4VG05WFUyUk5XVU5tU0dNeWJHUlBabXRoVGtWUk5USnphR3A2TnpnNWRrdFBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZLTUZKbkNsZHJOV2swV0RjeksyVkNjR2M1YVc4MGJuaFdkbGRWZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU5ha0pvVG1wRmVsbDZSWGRhUkd4cFRWZE5OVmt5U1RCTk1rVXhDazFYVlRST1ZFSnNUMWRhYUUxRVVUSlBWMFUxVG5wRk5VOVVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDFxUW1oT2FrVjZXWHBGZDFwRWJHbE5WMDAxV1RKSk1FMHlSVEZOVjFVMFRsUkNiRTlYV21nS1RVUlJNazlYUlRWT2VrVTFUMVJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDFxUW1oT2FrVjZXWHBGZDFwRWJHbE5WMDAxV1RKSk1FMHlSVEVLVFZkVk5FNVVRbXhQVjFwb1RVUlJNazlYUlRWT2VrVTFUMVJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUkpkMWxVV1hnS1RUSk5lRTFIVVRWWmFrWnFUMWRPYVU1RVRtaE9WRVpzVDBSVmQxcFViRzFaVkVFd1RtcHNhRTlVWTNoUFZHdDNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWtFd1RWUm5lVTVVU1RST2FtdDJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNDJNSFF6U3pSQlFVRlJSQXBCUldOM1VsRkphRUZPVEVkRGJERXhZa2xLUldWeFZsUkxiVk5ZUlhKTFVsUTJVelpRWkROcmVFZG9lRUV6YW04M01sQklRV2xCWlU4MGIzWnFiemRXQ21kdVdVeFVjek5XSzJsWE4zaEhURGg0Y2xOM2RHMU9URzE0TjNweGRIbG5jSHBCUzBKblozRm9hMnBQVUZGUlJFRjNUbTlCUkVKc1FXcEZRWEo2VFM4S1JETTJhMkl2YkhoTFlTdExZMloxVlVSaE9VNTBZbHB3VXpOUFlrZEVTall6WldSWVRpdDZVVmh1Vms1d1RHRnlheTlITkZsNFlYSnhaM1ExUVdwQ2FRcE9Nak0zVm1RMVdYUkxlRU5IS3pWWFpIZE1NMDF5TmtGeE9GVlNiekpFTkVkeGNsVkRLMWs0Vnk5dkswSkpWVWRqUWxCSE4weHhUV051UzJGYWRtODlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyDADAgEAMIICvwYJKoZIhvcNAQcCoIICsDCCAqwCAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgtP/dqEi51cKu2l85XfN5km5dNM0Ujr5itqTsd8DaAKwCFCbOe82yFVV1Y/Lyck9okXumYeiiGA8yMDI2MDcyOTA5MjEzMVowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdowggHWAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwNzI5MDkyMTMxWjAvBgkqhkiG9w0BCQQxIgQg5qRe4sooW/giyx5IFR5ECp988HljWzH2FAhreZRgyp8wgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGYwZAIwI1xw5cz+b1k4ZQzOqypDcrcG+lORfgw9lxHBEud5ftcXmyIvBVr4Pv35dRTf82NJAjA43A63oe13nximJ1xLxk4ySEQLPxEQN7+2/gg4lvmYBLzhfa4hUn+9u+hGmX8uXhE="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"HcRjUPLb3KuHEdA2qLMNrFDIMx3dbLFs0LmoGGUQUXo="}, "signature":"MEUCIQCYIevpqcNJhS3YAS+gryy9Fwikgx47hu0wXQYT8b+2gwIgIk2sJMKwUqjAZPJnLMMsr8cMyLxxu12xSUFaXimFJ90="}} \ No newline at end of file diff --git a/build/torch211-cxx11-cu130-x86_64-linux/modules.py b/build/torch211-cxx11-cu130-x86_64-linux/modules.py new file mode 100644 index 0000000000000000000000000000000000000000..222e384f2151c6a32275abcdb5000d5b8f1379d4 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/modules.py @@ -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, + ) diff --git a/build/torch211-cxx11-cu130-x86_64-linux/natten/__init__.py b/build/torch211-cxx11-cu130-x86_64-linux/natten/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/natten/__init__.py @@ -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"))) diff --git a/build/torch211-cxx11-cu130-x86_64-linux/token_permute/__init__.py b/build/torch211-cxx11-cu130-x86_64-linux/token_permute/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..381de37e4e7d4d2e6158668fbcc2ca52f7b388a3 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/token_permute/__init__.py @@ -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", +] diff --git a/build/torch211-cxx11-cu130-x86_64-linux/token_permute/cutlass_impl.py b/build/torch211-cxx11-cu130-x86_64-linux/token_permute/cutlass_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..49ce678f9052679448f843922f16bdec5d8796be --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/token_permute/cutlass_impl.py @@ -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 diff --git a/build/torch211-cxx11-cu130-x86_64-linux/token_permute/frontend.py b/build/torch211-cxx11-cu130-x86_64-linux/token_permute/frontend.py new file mode 100644 index 0000000000000000000000000000000000000000..720aded02e93e31f19a91efbda82cac3db15e4d2 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/token_permute/frontend.py @@ -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 diff --git a/build/torch211-cxx11-cu130-x86_64-linux/token_permute/torch_impl.py b/build/torch211-cxx11-cu130-x86_64-linux/token_permute/torch_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..7f06204e57cb6fc96ccb9e10ac9f798ac2d0e7a4 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/token_permute/torch_impl.py @@ -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", +] diff --git a/build/torch211-cxx11-cu130-x86_64-linux/utils/__init__.py b/build/torch211-cxx11-cu130-x86_64-linux/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..99f6c1e8cca79bfdf04640b8a92602b205a407a0 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/utils/__init__.py @@ -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. +# +################################################################################################# diff --git a/build/torch211-cxx11-cu130-x86_64-linux/utils/checks.py b/build/torch211-cxx11-cu130-x86_64-linux/utils/checks.py new file mode 100644 index 0000000000000000000000000000000000000000..f58db50a581e0bc6b43e8de0fbccd9404beaa1dc --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/utils/checks.py @@ -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, + ) diff --git a/build/torch211-cxx11-cu130-x86_64-linux/utils/device.py b/build/torch211-cxx11-cu130-x86_64-linux/utils/device.py new file mode 100644 index 0000000000000000000000000000000000000000..7997ab474197c78f3d41138f17ae311d6e06b1cf --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/utils/device.py @@ -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 diff --git a/build/torch211-cxx11-cu130-x86_64-linux/utils/dtype.py b/build/torch211-cxx11-cu130-x86_64-linux/utils/dtype.py new file mode 100644 index 0000000000000000000000000000000000000000..c65d9e9e0082c6f14af0b8331216cf26fcb3f28d --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/utils/dtype.py @@ -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] diff --git a/build/torch211-cxx11-cu130-x86_64-linux/utils/environment.py b/build/torch211-cxx11-cu130-x86_64-linux/utils/environment.py new file mode 100644 index 0000000000000000000000000000000000000000..0581aacb1183a5dea228a63859837113429e1a62 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/utils/environment.py @@ -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 diff --git a/build/torch211-cxx11-cu130-x86_64-linux/utils/log.py b/build/torch211-cxx11-cu130-x86_64-linux/utils/log.py new file mode 100644 index 0000000000000000000000000000000000000000..138c3ac8763480c0a381a8af0d913efcfa48b1a1 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/utils/log.py @@ -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) diff --git a/build/torch211-cxx11-cu130-x86_64-linux/utils/tensor.py b/build/torch211-cxx11-cu130-x86_64-linux/utils/tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..32d95acf2eb9a8c6fa29c7d3de3b48b5f14be415 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/utils/tensor.py @@ -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 diff --git a/build/torch211-cxx11-cu130-x86_64-linux/utils/testing.py b/build/torch211-cxx11-cu130-x86_64-linux/utils/testing.py new file mode 100644 index 0000000000000000000000000000000000000000..8ed9de5df24754dc9207a0b39954c8a73d7610a5 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/utils/testing.py @@ -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 diff --git a/build/torch211-cxx11-cu130-x86_64-linux/utils/tuples.py b/build/torch211-cxx11-cu130-x86_64-linux/utils/tuples.py new file mode 100644 index 0000000000000000000000000000000000000000..abf8a862fa4b647aab55fb2982e5dea08609d316 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/utils/tuples.py @@ -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 diff --git a/build/torch211-cxx11-cu130-x86_64-linux/utils/varlen.py b/build/torch211-cxx11-cu130-x86_64-linux/utils/varlen.py new file mode 100644 index 0000000000000000000000000000000000000000..8bc541dbd8e80ddf9923a097ac473b31035174ef --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/utils/varlen.py @@ -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, + ) diff --git a/build/torch211-cxx11-cu130-x86_64-linux/version.py b/build/torch211-cxx11-cu130-x86_64-linux/version.py new file mode 100644 index 0000000000000000000000000000000000000000..fe22ed72b87e6b242b01aae1ffb3de069ae7b0a9 --- /dev/null +++ b/build/torch211-cxx11-cu130-x86_64-linux/version.py @@ -0,0 +1,24 @@ +################################################################################################# +# 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. +# +################################################################################################# + +__version__ = "0.21.7" diff --git a/build/torch212-cxx11-cu126-x86_64-linux/__init__.py b/build/torch212-cxx11-cu126-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9fe10d400adaa8333f0da9c0c9654864659fbc3b --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/__init__.py @@ -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", +] diff --git a/build/torch212-cxx11-cu126-x86_64-linux/_environment.py b/build/torch212-cxx11-cu126-x86_64-linux/_environment.py new file mode 100644 index 0000000000000000000000000000000000000000..caae85bbe3c5eef992e4c6758e44b382382ebe41 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/_environment.py @@ -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", +] diff --git a/build/torch212-cxx11-cu126-x86_64-linux/_libnatten/__init__.py b/build/torch212-cxx11-cu126-x86_64-linux/_libnatten/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9566f8d4e7d9356415bb2951613dc5ada194dbd9 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/_libnatten/__init__.py @@ -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", +] diff --git a/build/torch212-cxx11-cu126-x86_64-linux/_libnatten/torch_wrappers.py b/build/torch212-cxx11-cu126-x86_64-linux/_libnatten/torch_wrappers.py new file mode 100644 index 0000000000000000000000000000000000000000..16e5b62b2bc5de0810f58d82d70d7fb2e3e56e4a --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/_libnatten/torch_wrappers.py @@ -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", +] diff --git a/build/torch212-cxx11-cu126-x86_64-linux/_natten_cuda_3641131.abi3.so b/build/torch212-cxx11-cu126-x86_64-linux/_natten_cuda_3641131.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..3039386b65a0192312f4f286d89f46d4e94fda5d --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/_natten_cuda_3641131.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bb631a10393b4c481cd6a3ed3972026ba991ab5157a8f6f34f122e385dce624b +size 103613616 diff --git a/build/torch212-cxx11-cu126-x86_64-linux/_ops.py b/build/torch212-cxx11-cu126-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..ee2d34306175408457b5e3dca950160c1444f947 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/_ops.py @@ -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}" diff --git a/build/torch212-cxx11-cu126-x86_64-linux/_types.py b/build/torch212-cxx11-cu126-x86_64-linux/_types.py new file mode 100644 index 0000000000000000000000000000000000000000..804980900b74586f5861f473853055c7d9630efd --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/_types.py @@ -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 diff --git a/build/torch212-cxx11-cu126-x86_64-linux/attn_merge.py b/build/torch212-cxx11-cu126-x86_64-linux/attn_merge.py new file mode 100644 index 0000000000000000000000000000000000000000..e71cc799d407df0afa701942ad08b03718ba20bf --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/attn_merge.py @@ -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"] diff --git a/build/torch212-cxx11-cu126-x86_64-linux/backends/__init__.py b/build/torch212-cxx11-cu126-x86_64-linux/backends/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..84b8192a37e048880aba2059a6f47e7d26f6a8ee --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/backends/__init__.py @@ -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", +] diff --git a/build/torch212-cxx11-cu126-x86_64-linux/backends/blackwell_fmha.py b/build/torch212-cxx11-cu126-x86_64-linux/backends/blackwell_fmha.py new file mode 100644 index 0000000000000000000000000000000000000000..9b36e79d8b9242b7c5a20833932721556abd61d5 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/backends/blackwell_fmha.py @@ -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 diff --git a/build/torch212-cxx11-cu126-x86_64-linux/backends/blackwell_fna.py b/build/torch212-cxx11-cu126-x86_64-linux/backends/blackwell_fna.py new file mode 100644 index 0000000000000000000000000000000000000000..5bd6f0ed8dd3da85bd60804e503f2aca27b37e4e --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/backends/blackwell_fna.py @@ -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, + ) diff --git a/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/__init__.py b/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9c2b75f7e5007dd1490e8ff5d4c6413f21a5e39f --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/__init__.py @@ -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) diff --git a/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/checks.py b/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/checks.py new file mode 100644 index 0000000000000000000000000000000000000000..4237d8eba4f2e11f7f7abfc796d53b87801f34a6 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/checks.py @@ -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 diff --git a/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass/__init__.py b/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fdd2d4e18d0d959b16f7c9bd0a8b2f71081845be --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass/__init__.py @@ -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}" + ) diff --git a/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass/backward_knobs.py b/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass/backward_knobs.py new file mode 100644 index 0000000000000000000000000000000000000000..dfb1c10a3f1dfa28597f6eb30e5f925ad4785325 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass/backward_knobs.py @@ -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.") diff --git a/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_backward_128x128.py b/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_backward_128x128.py new file mode 100644 index 0000000000000000000000000000000000000000..af6047ccb59f6923be82d6379b92c39b9bcb7f1a --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_backward_128x128.py @@ -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)), + ], +} diff --git a/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_backward_128x64.py b/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_backward_128x64.py new file mode 100644 index 0000000000000000000000000000000000000000..6a915e9686203919079092eccc743dce88b4107a --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_backward_128x64.py @@ -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)), + ], +} diff --git a/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_backward_64x64.py b/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_backward_64x64.py new file mode 100644 index 0000000000000000000000000000000000000000..8511bd844f34a263d233d4e6a1f6f7c775139e2d --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_backward_64x64.py @@ -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)), + ], +} diff --git a/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_forward_32x128.py b/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_forward_32x128.py new file mode 100644 index 0000000000000000000000000000000000000000..a930aba67df1de878d7e95797bb3b6d2d903e020 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_forward_32x128.py @@ -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)), + ], +} diff --git a/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_forward_64x128.py b/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_forward_64x128.py new file mode 100644 index 0000000000000000000000000000000000000000..1bc1fc78e028a755b755ea8108264eccf0bb9659 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_forward_64x128.py @@ -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)), + ], +} diff --git a/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_forward_64x64.py b/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_forward_64x64.py new file mode 100644 index 0000000000000000000000000000000000000000..c321c5b4e1a96f0cb58894f95dd2826d2d21c950 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_forward_64x64.py @@ -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)), + ], +} diff --git a/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass_blackwell/__init__.py b/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass_blackwell/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4a8c87c6c6d916f58c533fc2a84c66aff8bfe46b --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass_blackwell/__init__.py @@ -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}" + ) diff --git a/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass_hopper/__init__.py b/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass_hopper/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..70b2e3106003ed583050be356b54e32c9b7c61c0 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/cutlass_hopper/__init__.py @@ -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}" + ) diff --git a/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/flex/__init__.py b/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/flex/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c5027701152d3d5976670fd59dc158331c5cadfa --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/backends/configs/flex/__init__.py @@ -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}" + ) diff --git a/build/torch212-cxx11-cu126-x86_64-linux/backends/flex.py b/build/torch212-cxx11-cu126-x86_64-linux/backends/flex.py new file mode 100644 index 0000000000000000000000000000000000000000..0555a7ebdb62e6e4060fbfe3106bed116b44797b --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/backends/flex.py @@ -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, + ) diff --git a/build/torch212-cxx11-cu126-x86_64-linux/backends/fmha.py b/build/torch212-cxx11-cu126-x86_64-linux/backends/fmha.py new file mode 100644 index 0000000000000000000000000000000000000000..0c6851f8bb362a4c7d45db46e61db1ed1dc61d51 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/backends/fmha.py @@ -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 diff --git a/build/torch212-cxx11-cu126-x86_64-linux/backends/fna.py b/build/torch212-cxx11-cu126-x86_64-linux/backends/fna.py new file mode 100644 index 0000000000000000000000000000000000000000..747ecb1f7dcfe032e02e4fa76b8a3a9b70e96bd0 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/backends/fna.py @@ -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, + ) diff --git a/build/torch212-cxx11-cu126-x86_64-linux/backends/hopper_fmha.py b/build/torch212-cxx11-cu126-x86_64-linux/backends/hopper_fmha.py new file mode 100644 index 0000000000000000000000000000000000000000..251b50dc34598111f288ed3ac3d2b414ea23f253 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/backends/hopper_fmha.py @@ -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 diff --git a/build/torch212-cxx11-cu126-x86_64-linux/backends/hopper_fna.py b/build/torch212-cxx11-cu126-x86_64-linux/backends/hopper_fna.py new file mode 100644 index 0000000000000000000000000000000000000000..17bb6d5952675043b66b59fcf92ea9eea54facf9 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/backends/hopper_fna.py @@ -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, + ) diff --git a/build/torch212-cxx11-cu126-x86_64-linux/backends/reference.py b/build/torch212-cxx11-cu126-x86_64-linux/backends/reference.py new file mode 100644 index 0000000000000000000000000000000000000000..d0f34b41a8c9b1a1051360b23d7e20573931bb27 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/backends/reference.py @@ -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, + ) diff --git a/build/torch212-cxx11-cu126-x86_64-linux/context.py b/build/torch212-cxx11-cu126-x86_64-linux/context.py new file mode 100644 index 0000000000000000000000000000000000000000..1793ac984bbe5be383bec2c07b54c770c9969b42 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/context.py @@ -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) diff --git a/build/torch212-cxx11-cu126-x86_64-linux/functional.py b/build/torch212-cxx11-cu126-x86_64-linux/functional.py new file mode 100644 index 0000000000000000000000000000000000000000..944a0c2a7786de250fd5031e933133ea676d2a79 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/functional.py @@ -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, + ) diff --git a/build/torch212-cxx11-cu126-x86_64-linux/metadata.json b/build/torch212-cxx11-cu126-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..27dfa21f1234c168c956cdc9b628e0bc60812c43 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/metadata.json @@ -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": "u2MaEDk7TEgc1qPtOXICa6mRq1FXqPbzTxIuOF3OYks=", + "_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 + } + } +} \ No newline at end of file diff --git a/build/torch212-cxx11-cu126-x86_64-linux/metadata.json.sigstore b/build/torch212-cxx11-cu126-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..3891c8c0fcc0bbf1dcdec12641963d5c0cab827c --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHTDCCBtKgAwIBAgIUTZ9PFohpjt3exTwdQhHgbbQaqKAwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNzI5MDkyMTM1WhcNMjYwNzI5MDkzMTM1WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAED2MXl78JDYXF6Jv0F3K3imYBy4D/Z+n7jHm7hmHP3W/nE4mkejyRsPNgpOKhn9LWQhpXcliUagWhdMP3JVU/x6OCBfEwggXtMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUYlgYjU0AaBTpNIdwJSlDgpTtEJ8wHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoMjBhNjEzYzEwZDliMWM5Y2I0M2E1MWU4NTBlOWZhMDQ2OWE5NzE5OTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoMjBhNjEzYzEwZDliMWM5Y2I0M2E1MWU4NTBlOWZhMDQ2OWE5NzE5OTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoMjBhNjEzYzEwZDliMWM5Y2I0M2E1MWU4NTBlOWZhMDQ2OWE5NzE5OTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDIwYTYxM2MxMGQ5YjFjOWNiNDNhNTFlODUwZTlmYTA0NjlhOTcxOTkwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzA0MTgyNTI4NjkvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiwYKKwYBBAHWeQIEAgR9BHsAeQB3AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABn60t6WYAAAQDAEgwRgIhALK6DJt1ul8r6yhnsZ+F68OFu1esHLBVvys2nfJE+fB3AiEAroCnCyZuRIgiYAJpSTWpp211FnjCUD3jfQ2BLKQKLWIwCgYIKoZIzj0EAwMDaAAwZQIwd5BbmMzuH5hfiQYWZKL9iRGS29Y3gMlal5DJaeXOUPw4O4izLpwGOyemV71qWtaFAjEAgtdwuG4luQmDUFt7wIrSyq0K++C50qU6qtALbrk1Ge4zYW1FclAQxztS0fsOSLp0"}, "tlogEntries":[{"logIndex":"2280149242", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1785316895", "inclusionPromise":{"signedEntryTimestamp":"MEUCIQDSUuvjlqdxmQiB8jnQVaV2TKoTyqrSxN+HpH0AyJzk7wIgNMyEDACi+TKIRvl0HskByDKpH6DscrYKDEgLuQo21es="}, "inclusionProof":{"logIndex":"2158244980", "rootHash":"WF4VJ07vMLsN6d+3qqCLLT6FvaBbTn4Xn0BhaYi6aV8=", "treeSize":"2158244983", "hashes":["e9VaHKcKKjgWJi3kcwmC87A+isfuoT3OmLLLep71bhE=", "t+fABTSGdVomuhBIKNUD3jMsnQsKI7OKin1GdnblEWU=", "V/rjTAMZ2DuFyq7WqHqrrJGKGpmuUisyuONq9MtbuQE=", "3LrZZ5T93/7CInZzUe+S+IcEcc7ONboRD4Efht9iG/Y=", "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\n2158244983\nWF4VJ07vMLsN6d+3qqCLLT6FvaBbTn4Xn0BhaYi6aV8=\n\n— rekor.sigstore.dev wNI9ajBFAiBBa5AIbFxa34qe2KNNOfAUikNtbmPjwLt63tRSqCfSdgIhAIeoHWymWg0OvjxCsEKiL7SE6mQOFLqfEflGmoSwZ9Xc\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiIwNTFmNGFkOWQ4MGJhNmY4ZmNkY2U2YTdlYTU3NDc3NmEyMmQzNTdmYmFmMGEwMzMwZjVlNjk4ZDM5NzNjYzYzIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUNCb3BjUlRhdlJvekJSdVdGSVdhNEJybjNsU0lVYlBGdGs2WUhXUUpHUmVBSWdNQVBKamVta1d2ZFRVeER2MVRVVzdhMGNnM0QvbDFmSjhnZDhsQzFMbEc0PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFVSRU5EUW5STFowRjNTVUpCWjBsVlZGbzVVRVp2YUhCcWRETmxlRlIzWkZGb1NHZGlZbEZoY1V0QmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDU2U1RWTlJHdDVUVlJOTVZkb1kwNU5hbGwzVG5wSk5VMUVhM3BOVkUweFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZFTWsxWWJEYzRTa1JaV0VZMlNuWXdSak5MTTJsdFdVSjVORVF2V2l0dU4ycEliVGNLYUcxSVVETlhMMjVGTkcxclpXcDVVbk5RVG1kd1QwdG9iamxNVjFGb2NGaGpiR2xWWVdkWGFHUk5VRE5LVmxVdmVEWlBRMEptUlhkbloxaDBUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZaYkdkWkNtcFZNRUZoUWxSd1RrbGtkMHBUYkVSbmNGUjBSVW80ZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU5ha0pvVG1wRmVsbDZSWGRhUkd4cFRWZE5OVmt5U1RCTk1rVXhDazFYVlRST1ZFSnNUMWRhYUUxRVVUSlBWMFUxVG5wRk5VOVVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDFxUW1oT2FrVjZXWHBGZDFwRWJHbE5WMDAxV1RKSk1FMHlSVEZOVjFVMFRsUkNiRTlYV21nS1RVUlJNazlYUlRWT2VrVTFUMVJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDFxUW1oT2FrVjZXWHBGZDFwRWJHbE5WMDAxV1RKSk1FMHlSVEVLVFZkVk5FNVVRbXhQVjFwb1RVUlJNazlYUlRWT2VrVTFUMVJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUkpkMWxVV1hnS1RUSk5lRTFIVVRWWmFrWnFUMWRPYVU1RVRtaE9WRVpzVDBSVmQxcFViRzFaVkVFd1RtcHNhRTlVWTNoUFZHdDNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWtFd1RWUm5lVTVVU1RST2FtdDJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwZDFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT1VKSWMwRUtaVkZDTTBGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNDJNSFEyVjFsQlFVRlJSQXBCUldkM1VtZEphRUZNU3paRVNuUXhkV3c0Y2paNWFHNXpXaXRHTmpoUFJuVXhaWE5JVEVKV2RubHpNbTVtU2tVclprSXpRV2xGUVhKdlEyNURlVnAxQ2xKSloybFpRVXB3VTFSWGNIQXlNVEZHYm1wRFZVUXphbVpSTWtKTVMxRkxURmRKZDBObldVbExiMXBKZW1vd1JVRjNUVVJoUVVGM1dsRkpkMlExUW1JS2JVMTZkVWcxYUdacFVWbFhXa3RNT1dsU1IxTXlPVmt6WjAxc1lXdzFSRXBoWlZoUFZWQjNORTgwYVhwTWNIZEhUM2xsYlZZM01YRlhkR0ZHUVdwRlFRcG5kR1IzZFVjMGJIVlJiVVJWUm5RM2QwbHlVM2x4TUVzckswTTFNSEZWTm5GMFFVeGljbXN4UjJVMGVsbFhNVVpqYkVGUmVIcDBVekJtYzA5VFRIQXdDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgLM4zIUjfc9WEvmZX0LbkElqGWYJhT60DusjkMZIst5ECFAcNaueonkdZHo4PUCQGk9ieSGTgGA8yMDI2MDcyOTA5MjEzNVowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdswggHXAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwNzI5MDkyMTM1WjAvBgkqhkiG9w0BCQQxIgQgkfYh/92dEse6HTSCgmJh4lqAIIuYGso4uw6Qo4HLwIwwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGcwZQIxAMOw8btu8RRP3T+KphQU7NZBjJ7qaAy/5ta/Yu3QFh8KqOtDFY1149rMJoA2+z56nAIwXXLxg5Un4Ycx+th92YLSibTksWXqyT0392jEWBgP/yaICjkNlpT7CTcVWtX3nGP8"}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"BR9K2dgLpvj83Oan6ldHdqItNX+68KAzD15pjTlzzGM="}, "signature":"MEUCIQCBopcRTavRozBRuWFIWa4Brn3lSIUbPFtk6YHWQJGReAIgMAPJjemkWvdTUxDv1TUW7a0cg3D/l1fJ8gd8lC1LlG4="}} \ No newline at end of file diff --git a/build/torch212-cxx11-cu126-x86_64-linux/modules.py b/build/torch212-cxx11-cu126-x86_64-linux/modules.py new file mode 100644 index 0000000000000000000000000000000000000000..222e384f2151c6a32275abcdb5000d5b8f1379d4 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/modules.py @@ -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, + ) diff --git a/build/torch212-cxx11-cu126-x86_64-linux/natten/__init__.py b/build/torch212-cxx11-cu126-x86_64-linux/natten/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/natten/__init__.py @@ -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"))) diff --git a/build/torch212-cxx11-cu126-x86_64-linux/token_permute/__init__.py b/build/torch212-cxx11-cu126-x86_64-linux/token_permute/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..381de37e4e7d4d2e6158668fbcc2ca52f7b388a3 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/token_permute/__init__.py @@ -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", +] diff --git a/build/torch212-cxx11-cu126-x86_64-linux/token_permute/cutlass_impl.py b/build/torch212-cxx11-cu126-x86_64-linux/token_permute/cutlass_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..49ce678f9052679448f843922f16bdec5d8796be --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/token_permute/cutlass_impl.py @@ -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 diff --git a/build/torch212-cxx11-cu126-x86_64-linux/token_permute/frontend.py b/build/torch212-cxx11-cu126-x86_64-linux/token_permute/frontend.py new file mode 100644 index 0000000000000000000000000000000000000000..720aded02e93e31f19a91efbda82cac3db15e4d2 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/token_permute/frontend.py @@ -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 diff --git a/build/torch212-cxx11-cu126-x86_64-linux/token_permute/torch_impl.py b/build/torch212-cxx11-cu126-x86_64-linux/token_permute/torch_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..7f06204e57cb6fc96ccb9e10ac9f798ac2d0e7a4 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/token_permute/torch_impl.py @@ -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", +] diff --git a/build/torch212-cxx11-cu126-x86_64-linux/utils/__init__.py b/build/torch212-cxx11-cu126-x86_64-linux/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..99f6c1e8cca79bfdf04640b8a92602b205a407a0 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/utils/__init__.py @@ -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. +# +################################################################################################# diff --git a/build/torch212-cxx11-cu126-x86_64-linux/utils/checks.py b/build/torch212-cxx11-cu126-x86_64-linux/utils/checks.py new file mode 100644 index 0000000000000000000000000000000000000000..f58db50a581e0bc6b43e8de0fbccd9404beaa1dc --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/utils/checks.py @@ -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, + ) diff --git a/build/torch212-cxx11-cu126-x86_64-linux/utils/device.py b/build/torch212-cxx11-cu126-x86_64-linux/utils/device.py new file mode 100644 index 0000000000000000000000000000000000000000..7997ab474197c78f3d41138f17ae311d6e06b1cf --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/utils/device.py @@ -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 diff --git a/build/torch212-cxx11-cu126-x86_64-linux/utils/dtype.py b/build/torch212-cxx11-cu126-x86_64-linux/utils/dtype.py new file mode 100644 index 0000000000000000000000000000000000000000..c65d9e9e0082c6f14af0b8331216cf26fcb3f28d --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/utils/dtype.py @@ -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] diff --git a/build/torch212-cxx11-cu126-x86_64-linux/utils/environment.py b/build/torch212-cxx11-cu126-x86_64-linux/utils/environment.py new file mode 100644 index 0000000000000000000000000000000000000000..0581aacb1183a5dea228a63859837113429e1a62 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/utils/environment.py @@ -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 diff --git a/build/torch212-cxx11-cu126-x86_64-linux/utils/log.py b/build/torch212-cxx11-cu126-x86_64-linux/utils/log.py new file mode 100644 index 0000000000000000000000000000000000000000..138c3ac8763480c0a381a8af0d913efcfa48b1a1 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/utils/log.py @@ -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) diff --git a/build/torch212-cxx11-cu126-x86_64-linux/utils/tensor.py b/build/torch212-cxx11-cu126-x86_64-linux/utils/tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..32d95acf2eb9a8c6fa29c7d3de3b48b5f14be415 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/utils/tensor.py @@ -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 diff --git a/build/torch212-cxx11-cu126-x86_64-linux/utils/testing.py b/build/torch212-cxx11-cu126-x86_64-linux/utils/testing.py new file mode 100644 index 0000000000000000000000000000000000000000..8ed9de5df24754dc9207a0b39954c8a73d7610a5 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/utils/testing.py @@ -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 diff --git a/build/torch212-cxx11-cu126-x86_64-linux/utils/tuples.py b/build/torch212-cxx11-cu126-x86_64-linux/utils/tuples.py new file mode 100644 index 0000000000000000000000000000000000000000..abf8a862fa4b647aab55fb2982e5dea08609d316 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/utils/tuples.py @@ -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 diff --git a/build/torch212-cxx11-cu126-x86_64-linux/utils/varlen.py b/build/torch212-cxx11-cu126-x86_64-linux/utils/varlen.py new file mode 100644 index 0000000000000000000000000000000000000000..8bc541dbd8e80ddf9923a097ac473b31035174ef --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/utils/varlen.py @@ -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, + ) diff --git a/build/torch212-cxx11-cu126-x86_64-linux/version.py b/build/torch212-cxx11-cu126-x86_64-linux/version.py new file mode 100644 index 0000000000000000000000000000000000000000..fe22ed72b87e6b242b01aae1ffb3de069ae7b0a9 --- /dev/null +++ b/build/torch212-cxx11-cu126-x86_64-linux/version.py @@ -0,0 +1,24 @@ +################################################################################################# +# 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. +# +################################################################################################# + +__version__ = "0.21.7" diff --git a/build/torch212-cxx11-cu130-x86_64-linux/__init__.py b/build/torch212-cxx11-cu130-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9fe10d400adaa8333f0da9c0c9654864659fbc3b --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/__init__.py @@ -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", +] diff --git a/build/torch212-cxx11-cu130-x86_64-linux/_environment.py b/build/torch212-cxx11-cu130-x86_64-linux/_environment.py new file mode 100644 index 0000000000000000000000000000000000000000..caae85bbe3c5eef992e4c6758e44b382382ebe41 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/_environment.py @@ -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", +] diff --git a/build/torch212-cxx11-cu130-x86_64-linux/_libnatten/__init__.py b/build/torch212-cxx11-cu130-x86_64-linux/_libnatten/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9566f8d4e7d9356415bb2951613dc5ada194dbd9 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/_libnatten/__init__.py @@ -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", +] diff --git a/build/torch212-cxx11-cu130-x86_64-linux/_libnatten/torch_wrappers.py b/build/torch212-cxx11-cu130-x86_64-linux/_libnatten/torch_wrappers.py new file mode 100644 index 0000000000000000000000000000000000000000..16e5b62b2bc5de0810f58d82d70d7fb2e3e56e4a --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/_libnatten/torch_wrappers.py @@ -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", +] diff --git a/build/torch212-cxx11-cu130-x86_64-linux/_natten_cuda_3641131.abi3.so b/build/torch212-cxx11-cu130-x86_64-linux/_natten_cuda_3641131.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..886f1fc3dd19dbea726ce599db277fc334688e45 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/_natten_cuda_3641131.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e4d788cc59911b0c9687a966d37ada810872ddda4475de6a28089a76e8271724 +size 164387400 diff --git a/build/torch212-cxx11-cu130-x86_64-linux/_ops.py b/build/torch212-cxx11-cu130-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..ee2d34306175408457b5e3dca950160c1444f947 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/_ops.py @@ -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}" diff --git a/build/torch212-cxx11-cu130-x86_64-linux/_types.py b/build/torch212-cxx11-cu130-x86_64-linux/_types.py new file mode 100644 index 0000000000000000000000000000000000000000..804980900b74586f5861f473853055c7d9630efd --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/_types.py @@ -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 diff --git a/build/torch212-cxx11-cu130-x86_64-linux/attn_merge.py b/build/torch212-cxx11-cu130-x86_64-linux/attn_merge.py new file mode 100644 index 0000000000000000000000000000000000000000..e71cc799d407df0afa701942ad08b03718ba20bf --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/attn_merge.py @@ -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"] diff --git a/build/torch212-cxx11-cu130-x86_64-linux/backends/__init__.py b/build/torch212-cxx11-cu130-x86_64-linux/backends/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..84b8192a37e048880aba2059a6f47e7d26f6a8ee --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/backends/__init__.py @@ -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", +] diff --git a/build/torch212-cxx11-cu130-x86_64-linux/backends/blackwell_fmha.py b/build/torch212-cxx11-cu130-x86_64-linux/backends/blackwell_fmha.py new file mode 100644 index 0000000000000000000000000000000000000000..9b36e79d8b9242b7c5a20833932721556abd61d5 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/backends/blackwell_fmha.py @@ -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 diff --git a/build/torch212-cxx11-cu130-x86_64-linux/backends/blackwell_fna.py b/build/torch212-cxx11-cu130-x86_64-linux/backends/blackwell_fna.py new file mode 100644 index 0000000000000000000000000000000000000000..5bd6f0ed8dd3da85bd60804e503f2aca27b37e4e --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/backends/blackwell_fna.py @@ -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, + ) diff --git a/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/__init__.py b/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9c2b75f7e5007dd1490e8ff5d4c6413f21a5e39f --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/__init__.py @@ -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) diff --git a/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/checks.py b/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/checks.py new file mode 100644 index 0000000000000000000000000000000000000000..4237d8eba4f2e11f7f7abfc796d53b87801f34a6 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/checks.py @@ -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 diff --git a/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass/__init__.py b/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fdd2d4e18d0d959b16f7c9bd0a8b2f71081845be --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass/__init__.py @@ -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}" + ) diff --git a/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass/backward_knobs.py b/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass/backward_knobs.py new file mode 100644 index 0000000000000000000000000000000000000000..dfb1c10a3f1dfa28597f6eb30e5f925ad4785325 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass/backward_knobs.py @@ -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.") diff --git a/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_backward_128x128.py b/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_backward_128x128.py new file mode 100644 index 0000000000000000000000000000000000000000..af6047ccb59f6923be82d6379b92c39b9bcb7f1a --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_backward_128x128.py @@ -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)), + ], +} diff --git a/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_backward_128x64.py b/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_backward_128x64.py new file mode 100644 index 0000000000000000000000000000000000000000..6a915e9686203919079092eccc743dce88b4107a --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_backward_128x64.py @@ -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)), + ], +} diff --git a/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_backward_64x64.py b/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_backward_64x64.py new file mode 100644 index 0000000000000000000000000000000000000000..8511bd844f34a263d233d4e6a1f6f7c775139e2d --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_backward_64x64.py @@ -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)), + ], +} diff --git a/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_forward_32x128.py b/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_forward_32x128.py new file mode 100644 index 0000000000000000000000000000000000000000..a930aba67df1de878d7e95797bb3b6d2d903e020 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_forward_32x128.py @@ -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)), + ], +} diff --git a/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_forward_64x128.py b/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_forward_64x128.py new file mode 100644 index 0000000000000000000000000000000000000000..1bc1fc78e028a755b755ea8108264eccf0bb9659 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_forward_64x128.py @@ -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)), + ], +} diff --git a/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_forward_64x64.py b/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_forward_64x64.py new file mode 100644 index 0000000000000000000000000000000000000000..c321c5b4e1a96f0cb58894f95dd2826d2d21c950 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_forward_64x64.py @@ -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)), + ], +} diff --git a/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass_blackwell/__init__.py b/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass_blackwell/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4a8c87c6c6d916f58c533fc2a84c66aff8bfe46b --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass_blackwell/__init__.py @@ -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}" + ) diff --git a/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass_hopper/__init__.py b/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass_hopper/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..70b2e3106003ed583050be356b54e32c9b7c61c0 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/cutlass_hopper/__init__.py @@ -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}" + ) diff --git a/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/flex/__init__.py b/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/flex/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c5027701152d3d5976670fd59dc158331c5cadfa --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/backends/configs/flex/__init__.py @@ -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}" + ) diff --git a/build/torch212-cxx11-cu130-x86_64-linux/backends/flex.py b/build/torch212-cxx11-cu130-x86_64-linux/backends/flex.py new file mode 100644 index 0000000000000000000000000000000000000000..0555a7ebdb62e6e4060fbfe3106bed116b44797b --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/backends/flex.py @@ -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, + ) diff --git a/build/torch212-cxx11-cu130-x86_64-linux/backends/fmha.py b/build/torch212-cxx11-cu130-x86_64-linux/backends/fmha.py new file mode 100644 index 0000000000000000000000000000000000000000..0c6851f8bb362a4c7d45db46e61db1ed1dc61d51 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/backends/fmha.py @@ -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 diff --git a/build/torch212-cxx11-cu130-x86_64-linux/backends/fna.py b/build/torch212-cxx11-cu130-x86_64-linux/backends/fna.py new file mode 100644 index 0000000000000000000000000000000000000000..747ecb1f7dcfe032e02e4fa76b8a3a9b70e96bd0 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/backends/fna.py @@ -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, + ) diff --git a/build/torch212-cxx11-cu130-x86_64-linux/backends/hopper_fmha.py b/build/torch212-cxx11-cu130-x86_64-linux/backends/hopper_fmha.py new file mode 100644 index 0000000000000000000000000000000000000000..251b50dc34598111f288ed3ac3d2b414ea23f253 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/backends/hopper_fmha.py @@ -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 diff --git a/build/torch212-cxx11-cu130-x86_64-linux/backends/hopper_fna.py b/build/torch212-cxx11-cu130-x86_64-linux/backends/hopper_fna.py new file mode 100644 index 0000000000000000000000000000000000000000..17bb6d5952675043b66b59fcf92ea9eea54facf9 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/backends/hopper_fna.py @@ -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, + ) diff --git a/build/torch212-cxx11-cu130-x86_64-linux/backends/reference.py b/build/torch212-cxx11-cu130-x86_64-linux/backends/reference.py new file mode 100644 index 0000000000000000000000000000000000000000..d0f34b41a8c9b1a1051360b23d7e20573931bb27 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/backends/reference.py @@ -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, + ) diff --git a/build/torch212-cxx11-cu130-x86_64-linux/context.py b/build/torch212-cxx11-cu130-x86_64-linux/context.py new file mode 100644 index 0000000000000000000000000000000000000000..1793ac984bbe5be383bec2c07b54c770c9969b42 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/context.py @@ -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) diff --git a/build/torch212-cxx11-cu130-x86_64-linux/functional.py b/build/torch212-cxx11-cu130-x86_64-linux/functional.py new file mode 100644 index 0000000000000000000000000000000000000000..944a0c2a7786de250fd5031e933133ea676d2a79 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/functional.py @@ -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, + ) diff --git a/build/torch212-cxx11-cu130-x86_64-linux/metadata.json b/build/torch212-cxx11-cu130-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..bc1d4c47207bc8e2669d671065b8f9ae58374c62 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/metadata.json @@ -0,0 +1,84 @@ +{ + "name": "natten", + "id": "_natten_cuda_3641131", + "version": 1, + "license": "MIT", + "upstream": "https://github.com/SHI-Labs/NATTEN", + "python-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "10.0", + "10.0a", + "12.0", + "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": "5NeIzFmRGwyWh6lm03ragQhy3dpEdd5qKAiadugnFyQ=", + "_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 + } + } +} \ No newline at end of file diff --git a/build/torch212-cxx11-cu130-x86_64-linux/metadata.json.sigstore b/build/torch212-cxx11-cu130-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..5787709483ad6359e7f833fc03853d930be5e219 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtCgAwIBAgIUXWxIvUhwkLF53vZ/IoNNKlbKRtUwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNzI5MDkyMTM3WhcNMjYwNzI5MDkzMTM3WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEDcK0CMKZeT1K6XTIYynM21sU5WH1Hh/YU/wynI54bi8q4yrfVQH5YDh7yzNURm9/H4HI9waW18rWmeGy0Wyz3aOCBe8wggXrMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQU0m5WruquPybeZ8sXupU8LyIBMJYwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoMjBhNjEzYzEwZDliMWM5Y2I0M2E1MWU4NTBlOWZhMDQ2OWE5NzE5OTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoMjBhNjEzYzEwZDliMWM5Y2I0M2E1MWU4NTBlOWZhMDQ2OWE5NzE5OTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoMjBhNjEzYzEwZDliMWM5Y2I0M2E1MWU4NTBlOWZhMDQ2OWE5NzE5OTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDIwYTYxM2MxMGQ5YjFjOWNiNDNhNTFlODUwZTlmYTA0NjlhOTcxOTkwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzA0MTgyNTI4NjkvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiQYKKwYBBAHWeQIEAgR7BHkAdwB1AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABn60t8mgAAAQDAEYwRAIgSMzCwXB4e1Z45eXWM/PslYQQZ6YJis1B3aP2An1DZZUCIGTHmg9PGsQ+Wu2mh+S79VDswOVl5E6Sm8jgQQuHBPWyMAoGCCqGSM49BAMDA2kAMGYCMQD4InsM4Z4ZglNSXlgTQYRb38Ya+awycU5f7kLUkixWgDOWkfJ4GEusxYO1yMHANPACMQCBFDMNjA5zUcLElxWATJBI7wgZHqqff5PFt0+DIv1qIWdDjO7TzCDrFMQnMraEJi4="}, "tlogEntries":[{"logIndex":"2280149283", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1785316897", "inclusionPromise":{"signedEntryTimestamp":"MEUCIGu0GIj8ONKXRiULqDxO/ATNsVTRertlDO5VUu2yR1yHAiEA/OXwvAVfzCmkXLxXBa7yUP6Mo5RFk2lov79gHgEmv6w="}, "inclusionProof":{"logIndex":"2158245021", "rootHash":"8nsnkwCr/yklUz1Xs5gs5olFcpy8ld5qsZHUuE0rjAE=", "treeSize":"2158245025", "hashes":["li8Yh2oL/VqRV8jK8BDMHXO7wtvmJM2AVYLR47Tf7fw=", "DKC5pkSiTej/GdZHaZSWUL8eUiK5AlmLjT//RufQqnY=", "xGM/2PAvEOxv+7wOLWqCiT4P7PtVAhBgrBnnifQ+n80=", "wxdKCzS9dAOZvwsLQBJYZldyxOYJV7OAdiFwexwim8k=", "R62Rq5wJiW7Txz35eKsniX8PDDILzRh0W/n7Bwnhr54=", "6omgmCnbonyFyZqV6gto+LT4z9BdAvY8pslrc03aVb4=", "/5uu/jw8GTRTqEZ0auUV9lv+Zg0twg6hHe7ZChWeiJk=", "DwNueYXiyuVjob7c8CRNRN57ioovkJTOeOeOQIi6/pI=", "bXYdybsl/2BK7Yled3FQJEE3Hk0TeGcWnXUxOkVZv8w=", "xrwx0yHkH5ZmCgeWs/jLZZ8RcejqpeuUeAWNJkAYrHk=", "8MN6j2GHUjPnNtIwpas9l/XjdwE0/KPUhFDlrVkngko=", "b+xUZfuENQxvSOJxzNvYvRG8eVphfszPpZmuf4/cQ6c=", "OVsvZCKnWA+498QUIaQCtitUT6huDbC7SmhH1l8MxXI=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2158245025\n8nsnkwCr/yklUz1Xs5gs5olFcpy8ld5qsZHUuE0rjAE=\n\n— rekor.sigstore.dev wNI9ajBGAiEA7KxCJuFeV8GeMSBnIi4N6xSRJAk4nqoH2etazXnDReQCIQCulVJcHXcE+HSbO/0V0TVbk1y5SUtJV0wiOp8Trois+g==\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJmMWNlOGYwN2MwOWM3OTAzOGE2NGNlMTE0ZThmMWRhNDU2ODIwMDEzNTAwZDNjZDkyZjFkNWMxNTNiNmIzZDEzIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUR4SktMNU0vOE5rYlg3dmhJS0wwdkNCam5KMm1XaGxDQUJoSW5jQTlFTmNnSWdEUGhmY3IwM0Y5MStmUXJrUmw0eW1OYy9ER0VBa045OUlvaWNUanptejlBPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5SRFowRjNTVUpCWjBsVldGZDRTWFpWYUhkclRFWTFNM1phTDBsdlRrNUxiR0pMVW5SVmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDU2U1RWTlJHdDVUVlJOTTFkb1kwNU5hbGwzVG5wSk5VMUVhM3BOVkUwelYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZFWTBzd1EwMUxXbVZVTVVzMldGUkpXWGx1VFRJeGMxVTFWMGd4U0dndldWVXZkM2tLYmtrMU5HSnBPSEUwZVhKbVZsRklOVmxFYURkNWVrNVZVbTA1TDBnMFNFazVkMkZYTVRoeVYyMWxSM2t3VjNsNk0yRlBRMEpsT0hkbloxaHlUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlV3YlRWWENuSjFjWFZRZVdKbFdqaHpXSFZ3VlRoTWVVbENUVXBaZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU5ha0pvVG1wRmVsbDZSWGRhUkd4cFRWZE5OVmt5U1RCTk1rVXhDazFYVlRST1ZFSnNUMWRhYUUxRVVUSlBWMFUxVG5wRk5VOVVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDFxUW1oT2FrVjZXWHBGZDFwRWJHbE5WMDAxV1RKSk1FMHlSVEZOVjFVMFRsUkNiRTlYV21nS1RVUlJNazlYUlRWT2VrVTFUMVJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDFxUW1oT2FrVjZXWHBGZDFwRWJHbE5WMDAxV1RKSk1FMHlSVEVLVFZkVk5FNVVRbXhQVjFwb1RVUlJNazlYUlRWT2VrVTFUMVJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUkpkMWxVV1hnS1RUSk5lRTFIVVRWWmFrWnFUMWRPYVU1RVRtaE9WRVpzVDBSVmQxcFViRzFaVkVFd1RtcHNhRTlVWTNoUFZHdDNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWtFd1RWUm5lVTVVU1RST2FtdDJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwVVZsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTTjBKSWEwRUtaSGRDTVVGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNDJNSFE0YldkQlFVRlJSQXBCUlZsM1VrRkpaMU5OZWtOM1dFSTBaVEZhTkRWbFdGZE5MMUJ6YkZsUlVWbzJXVXBwY3pGQ00yRlFNa0Z1TVVSYVdsVkRTVWRVU0cxbk9WQkhjMUVyQ2xkMU1tMW9LMU0zT1ZaRWMzZFBWbXcxUlRaVGJUaHFaMUZSZFVoQ1VGZDVUVUZ2UjBORGNVZFRUVFE1UWtGTlJFRXlhMEZOUjFsRFRWRkVORWx1YzAwS05GbzBXbWRzVGxOWWJHZFVVVmxTWWpNNFdXRXJZWGQ1WTFVMVpqZHJURlZyYVhoWFowUlBWMnRtU2pSSFJYVnplRmxQTVhsTlNFRk9VRUZEVFZGRFFncEdSRTFPYWtFMWVsVmpURVZzZUZkQlZFcENTVGQzWjFwSWNYRm1aalZRUm5Rd0swUkpkakZ4U1Zka1JHcFBOMVI2UTBSeVJrMVJiazF5WVVWS2FUUTlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgXvAJ9uabDITc3/OVPR/J/oFKNsCd9EtPs7sB776oHBcCFQCVVfG//WsGc2NlB7dOzsQ8SqQG8hgPMjAyNjA3MjkwOTIxMzdaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHaMIIB1gIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDcyOTA5MjEzN1owLwYJKoZIhvcNAQkEMSIEIM3aOx0YLolIsgZqPeimTbi6sCQ0ntuWJgeEedvL83kpMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRmMGQCMEITSFpfYav2h9BWWNKJf+Q7sQYoqz4u08uHaFpyX3VOFRtXchR0EDhh/2xXY56SdwIwTaYUz9bum5jXpnGalG0munckTLT6zOqVdHhHHjp2HH0IMb4h5nfyZNF9CXoZswJ2"}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"8c6PB8CceQOKZM4RTo8dpFaCABNQDTzZLx1cFTtrPRM="}, "signature":"MEUCIQDxJKL5M/8NkbX7vhIKL0vCBjnJ2mWhlCABhIncA9ENcgIgDPhfcr03F91+fQrkRl4ymNc/DGEAkN99IoicTjzmz9A="}} \ No newline at end of file diff --git a/build/torch212-cxx11-cu130-x86_64-linux/modules.py b/build/torch212-cxx11-cu130-x86_64-linux/modules.py new file mode 100644 index 0000000000000000000000000000000000000000..222e384f2151c6a32275abcdb5000d5b8f1379d4 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/modules.py @@ -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, + ) diff --git a/build/torch212-cxx11-cu130-x86_64-linux/natten/__init__.py b/build/torch212-cxx11-cu130-x86_64-linux/natten/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/natten/__init__.py @@ -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"))) diff --git a/build/torch212-cxx11-cu130-x86_64-linux/token_permute/__init__.py b/build/torch212-cxx11-cu130-x86_64-linux/token_permute/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..381de37e4e7d4d2e6158668fbcc2ca52f7b388a3 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/token_permute/__init__.py @@ -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", +] diff --git a/build/torch212-cxx11-cu130-x86_64-linux/token_permute/cutlass_impl.py b/build/torch212-cxx11-cu130-x86_64-linux/token_permute/cutlass_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..49ce678f9052679448f843922f16bdec5d8796be --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/token_permute/cutlass_impl.py @@ -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 diff --git a/build/torch212-cxx11-cu130-x86_64-linux/token_permute/frontend.py b/build/torch212-cxx11-cu130-x86_64-linux/token_permute/frontend.py new file mode 100644 index 0000000000000000000000000000000000000000..720aded02e93e31f19a91efbda82cac3db15e4d2 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/token_permute/frontend.py @@ -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 diff --git a/build/torch212-cxx11-cu130-x86_64-linux/token_permute/torch_impl.py b/build/torch212-cxx11-cu130-x86_64-linux/token_permute/torch_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..7f06204e57cb6fc96ccb9e10ac9f798ac2d0e7a4 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/token_permute/torch_impl.py @@ -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", +] diff --git a/build/torch212-cxx11-cu130-x86_64-linux/utils/__init__.py b/build/torch212-cxx11-cu130-x86_64-linux/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..99f6c1e8cca79bfdf04640b8a92602b205a407a0 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/utils/__init__.py @@ -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. +# +################################################################################################# diff --git a/build/torch212-cxx11-cu130-x86_64-linux/utils/checks.py b/build/torch212-cxx11-cu130-x86_64-linux/utils/checks.py new file mode 100644 index 0000000000000000000000000000000000000000..f58db50a581e0bc6b43e8de0fbccd9404beaa1dc --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/utils/checks.py @@ -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, + ) diff --git a/build/torch212-cxx11-cu130-x86_64-linux/utils/device.py b/build/torch212-cxx11-cu130-x86_64-linux/utils/device.py new file mode 100644 index 0000000000000000000000000000000000000000..7997ab474197c78f3d41138f17ae311d6e06b1cf --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/utils/device.py @@ -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 diff --git a/build/torch212-cxx11-cu130-x86_64-linux/utils/dtype.py b/build/torch212-cxx11-cu130-x86_64-linux/utils/dtype.py new file mode 100644 index 0000000000000000000000000000000000000000..c65d9e9e0082c6f14af0b8331216cf26fcb3f28d --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/utils/dtype.py @@ -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] diff --git a/build/torch212-cxx11-cu130-x86_64-linux/utils/environment.py b/build/torch212-cxx11-cu130-x86_64-linux/utils/environment.py new file mode 100644 index 0000000000000000000000000000000000000000..0581aacb1183a5dea228a63859837113429e1a62 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/utils/environment.py @@ -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 diff --git a/build/torch212-cxx11-cu130-x86_64-linux/utils/log.py b/build/torch212-cxx11-cu130-x86_64-linux/utils/log.py new file mode 100644 index 0000000000000000000000000000000000000000..138c3ac8763480c0a381a8af0d913efcfa48b1a1 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/utils/log.py @@ -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) diff --git a/build/torch212-cxx11-cu130-x86_64-linux/utils/tensor.py b/build/torch212-cxx11-cu130-x86_64-linux/utils/tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..32d95acf2eb9a8c6fa29c7d3de3b48b5f14be415 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/utils/tensor.py @@ -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 diff --git a/build/torch212-cxx11-cu130-x86_64-linux/utils/testing.py b/build/torch212-cxx11-cu130-x86_64-linux/utils/testing.py new file mode 100644 index 0000000000000000000000000000000000000000..8ed9de5df24754dc9207a0b39954c8a73d7610a5 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/utils/testing.py @@ -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 diff --git a/build/torch212-cxx11-cu130-x86_64-linux/utils/tuples.py b/build/torch212-cxx11-cu130-x86_64-linux/utils/tuples.py new file mode 100644 index 0000000000000000000000000000000000000000..abf8a862fa4b647aab55fb2982e5dea08609d316 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/utils/tuples.py @@ -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 diff --git a/build/torch212-cxx11-cu130-x86_64-linux/utils/varlen.py b/build/torch212-cxx11-cu130-x86_64-linux/utils/varlen.py new file mode 100644 index 0000000000000000000000000000000000000000..8bc541dbd8e80ddf9923a097ac473b31035174ef --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/utils/varlen.py @@ -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, + ) diff --git a/build/torch212-cxx11-cu130-x86_64-linux/version.py b/build/torch212-cxx11-cu130-x86_64-linux/version.py new file mode 100644 index 0000000000000000000000000000000000000000..fe22ed72b87e6b242b01aae1ffb3de069ae7b0a9 --- /dev/null +++ b/build/torch212-cxx11-cu130-x86_64-linux/version.py @@ -0,0 +1,24 @@ +################################################################################################# +# 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. +# +################################################################################################# + +__version__ = "0.21.7" diff --git a/build/torch212-cxx11-cu132-x86_64-linux/__init__.py b/build/torch212-cxx11-cu132-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9fe10d400adaa8333f0da9c0c9654864659fbc3b --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/__init__.py @@ -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", +] diff --git a/build/torch212-cxx11-cu132-x86_64-linux/_environment.py b/build/torch212-cxx11-cu132-x86_64-linux/_environment.py new file mode 100644 index 0000000000000000000000000000000000000000..caae85bbe3c5eef992e4c6758e44b382382ebe41 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/_environment.py @@ -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", +] diff --git a/build/torch212-cxx11-cu132-x86_64-linux/_libnatten/__init__.py b/build/torch212-cxx11-cu132-x86_64-linux/_libnatten/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9566f8d4e7d9356415bb2951613dc5ada194dbd9 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/_libnatten/__init__.py @@ -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", +] diff --git a/build/torch212-cxx11-cu132-x86_64-linux/_libnatten/torch_wrappers.py b/build/torch212-cxx11-cu132-x86_64-linux/_libnatten/torch_wrappers.py new file mode 100644 index 0000000000000000000000000000000000000000..16e5b62b2bc5de0810f58d82d70d7fb2e3e56e4a --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/_libnatten/torch_wrappers.py @@ -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", +] diff --git a/build/torch212-cxx11-cu132-x86_64-linux/_natten_cuda_3641131.abi3.so b/build/torch212-cxx11-cu132-x86_64-linux/_natten_cuda_3641131.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..cd1e665bbf35f2cb595a7b09174d56915923c7c4 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/_natten_cuda_3641131.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:808beb6538e39beb94c2937ee8a5fade8c1edbf60f08e82a97c83047c9d1c4d8 +size 166058392 diff --git a/build/torch212-cxx11-cu132-x86_64-linux/_ops.py b/build/torch212-cxx11-cu132-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..ee2d34306175408457b5e3dca950160c1444f947 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/_ops.py @@ -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}" diff --git a/build/torch212-cxx11-cu132-x86_64-linux/_types.py b/build/torch212-cxx11-cu132-x86_64-linux/_types.py new file mode 100644 index 0000000000000000000000000000000000000000..804980900b74586f5861f473853055c7d9630efd --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/_types.py @@ -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 diff --git a/build/torch212-cxx11-cu132-x86_64-linux/attn_merge.py b/build/torch212-cxx11-cu132-x86_64-linux/attn_merge.py new file mode 100644 index 0000000000000000000000000000000000000000..e71cc799d407df0afa701942ad08b03718ba20bf --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/attn_merge.py @@ -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"] diff --git a/build/torch212-cxx11-cu132-x86_64-linux/backends/__init__.py b/build/torch212-cxx11-cu132-x86_64-linux/backends/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..84b8192a37e048880aba2059a6f47e7d26f6a8ee --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/backends/__init__.py @@ -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", +] diff --git a/build/torch212-cxx11-cu132-x86_64-linux/backends/blackwell_fmha.py b/build/torch212-cxx11-cu132-x86_64-linux/backends/blackwell_fmha.py new file mode 100644 index 0000000000000000000000000000000000000000..9b36e79d8b9242b7c5a20833932721556abd61d5 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/backends/blackwell_fmha.py @@ -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 diff --git a/build/torch212-cxx11-cu132-x86_64-linux/backends/blackwell_fna.py b/build/torch212-cxx11-cu132-x86_64-linux/backends/blackwell_fna.py new file mode 100644 index 0000000000000000000000000000000000000000..5bd6f0ed8dd3da85bd60804e503f2aca27b37e4e --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/backends/blackwell_fna.py @@ -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, + ) diff --git a/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/__init__.py b/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9c2b75f7e5007dd1490e8ff5d4c6413f21a5e39f --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/__init__.py @@ -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) diff --git a/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/checks.py b/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/checks.py new file mode 100644 index 0000000000000000000000000000000000000000..4237d8eba4f2e11f7f7abfc796d53b87801f34a6 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/checks.py @@ -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 diff --git a/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass/__init__.py b/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fdd2d4e18d0d959b16f7c9bd0a8b2f71081845be --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass/__init__.py @@ -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}" + ) diff --git a/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass/backward_knobs.py b/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass/backward_knobs.py new file mode 100644 index 0000000000000000000000000000000000000000..dfb1c10a3f1dfa28597f6eb30e5f925ad4785325 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass/backward_knobs.py @@ -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.") diff --git a/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_backward_128x128.py b/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_backward_128x128.py new file mode 100644 index 0000000000000000000000000000000000000000..af6047ccb59f6923be82d6379b92c39b9bcb7f1a --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_backward_128x128.py @@ -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)), + ], +} diff --git a/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_backward_128x64.py b/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_backward_128x64.py new file mode 100644 index 0000000000000000000000000000000000000000..6a915e9686203919079092eccc743dce88b4107a --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_backward_128x64.py @@ -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)), + ], +} diff --git a/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_backward_64x64.py b/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_backward_64x64.py new file mode 100644 index 0000000000000000000000000000000000000000..8511bd844f34a263d233d4e6a1f6f7c775139e2d --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_backward_64x64.py @@ -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)), + ], +} diff --git a/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_forward_32x128.py b/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_forward_32x128.py new file mode 100644 index 0000000000000000000000000000000000000000..a930aba67df1de878d7e95797bb3b6d2d903e020 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_forward_32x128.py @@ -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)), + ], +} diff --git a/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_forward_64x128.py b/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_forward_64x128.py new file mode 100644 index 0000000000000000000000000000000000000000..1bc1fc78e028a755b755ea8108264eccf0bb9659 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_forward_64x128.py @@ -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)), + ], +} diff --git a/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_forward_64x64.py b/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_forward_64x64.py new file mode 100644 index 0000000000000000000000000000000000000000..c321c5b4e1a96f0cb58894f95dd2826d2d21c950 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_forward_64x64.py @@ -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)), + ], +} diff --git a/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass_blackwell/__init__.py b/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass_blackwell/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4a8c87c6c6d916f58c533fc2a84c66aff8bfe46b --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass_blackwell/__init__.py @@ -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}" + ) diff --git a/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass_hopper/__init__.py b/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass_hopper/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..70b2e3106003ed583050be356b54e32c9b7c61c0 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/cutlass_hopper/__init__.py @@ -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}" + ) diff --git a/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/flex/__init__.py b/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/flex/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c5027701152d3d5976670fd59dc158331c5cadfa --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/backends/configs/flex/__init__.py @@ -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}" + ) diff --git a/build/torch212-cxx11-cu132-x86_64-linux/backends/flex.py b/build/torch212-cxx11-cu132-x86_64-linux/backends/flex.py new file mode 100644 index 0000000000000000000000000000000000000000..0555a7ebdb62e6e4060fbfe3106bed116b44797b --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/backends/flex.py @@ -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, + ) diff --git a/build/torch212-cxx11-cu132-x86_64-linux/backends/fmha.py b/build/torch212-cxx11-cu132-x86_64-linux/backends/fmha.py new file mode 100644 index 0000000000000000000000000000000000000000..0c6851f8bb362a4c7d45db46e61db1ed1dc61d51 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/backends/fmha.py @@ -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 diff --git a/build/torch212-cxx11-cu132-x86_64-linux/backends/fna.py b/build/torch212-cxx11-cu132-x86_64-linux/backends/fna.py new file mode 100644 index 0000000000000000000000000000000000000000..747ecb1f7dcfe032e02e4fa76b8a3a9b70e96bd0 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/backends/fna.py @@ -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, + ) diff --git a/build/torch212-cxx11-cu132-x86_64-linux/backends/hopper_fmha.py b/build/torch212-cxx11-cu132-x86_64-linux/backends/hopper_fmha.py new file mode 100644 index 0000000000000000000000000000000000000000..251b50dc34598111f288ed3ac3d2b414ea23f253 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/backends/hopper_fmha.py @@ -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 diff --git a/build/torch212-cxx11-cu132-x86_64-linux/backends/hopper_fna.py b/build/torch212-cxx11-cu132-x86_64-linux/backends/hopper_fna.py new file mode 100644 index 0000000000000000000000000000000000000000..17bb6d5952675043b66b59fcf92ea9eea54facf9 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/backends/hopper_fna.py @@ -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, + ) diff --git a/build/torch212-cxx11-cu132-x86_64-linux/backends/reference.py b/build/torch212-cxx11-cu132-x86_64-linux/backends/reference.py new file mode 100644 index 0000000000000000000000000000000000000000..d0f34b41a8c9b1a1051360b23d7e20573931bb27 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/backends/reference.py @@ -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, + ) diff --git a/build/torch212-cxx11-cu132-x86_64-linux/context.py b/build/torch212-cxx11-cu132-x86_64-linux/context.py new file mode 100644 index 0000000000000000000000000000000000000000..1793ac984bbe5be383bec2c07b54c770c9969b42 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/context.py @@ -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) diff --git a/build/torch212-cxx11-cu132-x86_64-linux/functional.py b/build/torch212-cxx11-cu132-x86_64-linux/functional.py new file mode 100644 index 0000000000000000000000000000000000000000..944a0c2a7786de250fd5031e933133ea676d2a79 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/functional.py @@ -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, + ) diff --git a/build/torch212-cxx11-cu132-x86_64-linux/metadata.json b/build/torch212-cxx11-cu132-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..ac63e64d42c62dfe1efcf092b8b07cae40d1ed0a --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/metadata.json @@ -0,0 +1,84 @@ +{ + "name": "natten", + "id": "_natten_cuda_3641131", + "version": 1, + "license": "MIT", + "upstream": "https://github.com/SHI-Labs/NATTEN", + "python-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "10.0", + "10.0a", + "12.0", + "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": "gIvrZTjjm+uUwpN+6KX63owe2/YPCOgql8gwR8nRxNg=", + "_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 + } + } +} \ No newline at end of file diff --git a/build/torch212-cxx11-cu132-x86_64-linux/metadata.json.sigstore b/build/torch212-cxx11-cu132-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..27c3355464f4ae8ae785889b47ecec166326c90e --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSjCCBtGgAwIBAgIUdA6+XeYS+sRVZnmC578YXv7+c+wwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNzI5MDkyMTMyWhcNMjYwNzI5MDkzMTMyWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEEf8BXV7aO8ean0ijY9S0JIxGke7rUtAFoMCoFi+OFLeB8dlixmCscKBnRQgt/FlobtnDdzZDQ6ET/x3BYj9lWaOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUhrJBKBLQKHyJUoAt4/PCHCLmqz4wHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoMjBhNjEzYzEwZDliMWM5Y2I0M2E1MWU4NTBlOWZhMDQ2OWE5NzE5OTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoMjBhNjEzYzEwZDliMWM5Y2I0M2E1MWU4NTBlOWZhMDQ2OWE5NzE5OTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoMjBhNjEzYzEwZDliMWM5Y2I0M2E1MWU4NTBlOWZhMDQ2OWE5NzE5OTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDIwYTYxM2MxMGQ5YjFjOWNiNDNhNTFlODUwZTlmYTA0NjlhOTcxOTkwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzA0MTgyNTI4NjkvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABn60t4E4AAAQDAEcwRQIhAPSLAjfojH3mBfNFE+D7Hu1B75XpxzTIjsHxnIRpgwmqAiAAkFIG3LY0W/AgySrOTCIO8HxgwpfgtDLALjo9bnGSMzAKBggqhkjOPQQDAwNnADBkAjBY12Pv/xy8bSQ4yBMhB92ITpwM9diQ4XS2ukrPpP41keKf0rOWb/jfB38cFrSv8rECMFw+YGhuyrIvdbtJX3MSI6gwyzxnRj949E58DXEiIWjKCDIbQHaK6an7EypjXyKDZQ=="}, "tlogEntries":[{"logIndex":"2280149204", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1785316892", "inclusionPromise":{"signedEntryTimestamp":"MEYCIQDnNwUaG/8oJaotllmsRrWaD4C8eSZiAS6PrkWCmDpNSQIhAOdsbCuWBDoT9wu+4/LV6QEcOND1rg2zKR9bA2Hv4z95"}, "inclusionProof":{"logIndex":"2158244942", "rootHash":"SoLUrEf4SX0KKmvbCC0juc9GhCmAqQceojyZKcU1gdA=", "treeSize":"2158244944", "hashes":["TMzXFuI3uYKRdj9zDfW5Je0SuKbx9OGwLGNlIFWKR3g=", "fKE6YFdwRs4t7A6qL585NUbF53TTdlzSrFzXGIZ3LJQ=", "74wQqvWNc9nlETBtxT1sCnvA1hQnpX3lujr/ZJuwWGE=", "j22TWQ24Q2Vd5MvrakcXQADgmxNyt1e9AT4i8q6IrP0=", "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\n2158244944\nSoLUrEf4SX0KKmvbCC0juc9GhCmAqQceojyZKcU1gdA=\n\n— rekor.sigstore.dev wNI9ajBFAiAWjkYWbuKY57pZlwTAcqJ0PhJghtGCd4KUQii9LvPlkgIhAP/MTs+Qy8JryHRL0jkdk8LEK37hQJZAQJZ6FNs9jN/i\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiIxZGU0NjJkYjk0YzRiZDEzYzc0NjE3NjU5ODRjMDA1NmEzZjliZWI0YjM3MGI5MTI1NTVlZjRhYzJhMWUxOTgwIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUNUQXFuYTF6eS9SWGplcGRyYzd0aktWYmd6OTYydVRyN0dzMEx4dkYzMElnSWdLU0hRRzBmNmRCaVhzVDBJa1pYNUg0d2x1K1p2QWxUQXVGZGxpQ0h2YkFVPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRha05EUW5SSFowRjNTVUpCWjBsVlpFRTJLMWhsV1ZNcmMxSldXbTV0UXpVM09GbFlkamNyWXl0M2QwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDU2U1RWTlJHdDVUVlJOZVZkb1kwNU5hbGwzVG5wSk5VMUVhM3BOVkUxNVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZGWmpoQ1dGWTNZVTg0WldGdU1HbHFXVGxUTUVwSmVFZHJaVGR5VlhSQlJtOU5RMjhLUm1rclQwWk1aVUk0Wkd4cGVHMURjMk5MUW01U1VXZDBMMFpzYjJKMGJrUmtlbHBFVVRaRlZDOTRNMEpaYWpsc1YyRlBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZvY2twQ0NrdENURkZMU0hsS1ZXOUJkRFF2VUVOSVEweHRjWG8wZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU5ha0pvVG1wRmVsbDZSWGRhUkd4cFRWZE5OVmt5U1RCTk1rVXhDazFYVlRST1ZFSnNUMWRhYUUxRVVUSlBWMFUxVG5wRk5VOVVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDFxUW1oT2FrVjZXWHBGZDFwRWJHbE5WMDAxV1RKSk1FMHlSVEZOVjFVMFRsUkNiRTlYV21nS1RVUlJNazlYUlRWT2VrVTFUMVJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDFxUW1oT2FrVjZXWHBGZDFwRWJHbE5WMDAxV1RKSk1FMHlSVEVLVFZkVk5FNVVRbXhQVjFwb1RVUlJNazlYUlRWT2VrVTFUMVJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUkpkMWxVV1hnS1RUSk5lRTFIVVRWWmFrWnFUMWRPYVU1RVRtaE9WRVpzVDBSVmQxcFViRzFaVkVFd1RtcHNhRTlVWTNoUFZHdDNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWtFd1RWUm5lVTVVU1RST2FtdDJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNDJNSFEwUlRSQlFVRlJSQXBCUldOM1VsRkphRUZRVTB4QmFtWnZha2d6YlVKbVRrWkZLMFEzU0hVeFFqYzFXSEI0ZWxSSmFuTkllRzVKVW5CbmQyMXhRV2xCUVd0R1NVY3pURmt3Q2xjdlFXZDVVM0pQVkVOSlR6aEllR2QzY0dabmRFUk1RVXhxYnpsaWJrZFRUWHBCUzBKblozRm9hMnBQVUZGUlJFRjNUbTVCUkVKclFXcENXVEV5VUhZS0wzaDVPR0pUVVRSNVFrMW9Ramt5U1ZSd2QwMDVaR2xSTkZoVE1uVnJjbEJ3VURReGEyVkxaakJ5VDFkaUwycG1Rak00WTBaeVUzWTRja1ZEVFVaM0t3cFpSMmgxZVhKSmRtUmlkRXBZTTAxVFNUWm5kM2w2ZUc1U2FqazBPVVUxT0VSWVJXbEpWMnBMUTBSSllsRklZVXMyWVc0M1JYbHdhbGg1UzBSYVVUMDlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyzADAgEAMIICwgYJKoZIhvcNAQcCoIICszCCAq8CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgNQsJsvFeN4A8iudM5XNrFbSHC2juGVO5mkZ0uyyVsQMCFQDXOxklKPfDDJIRFkaaWguQpzCnSBgPMjAyNjA3MjkwOTIxMzJaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHcMIIB2AIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDcyOTA5MjEzMlowLwYJKoZIhvcNAQkEMSIEIAorQDJ+Gmd/2G8poo/eTjOyRba5yWSqE8w2HNowHvRGMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRoMGYCMQDMoiOnBLXbSB0QQeyARr0sm2P/a3hDiLzhN0uJVVDp4u8NTpNm+JMwOeWT4JAha+cCMQDn/cy6/iY87AC3FA8ZNKuSxr4FVyhJJA+vNQwrI2xa+xYCMBaM1dyVoxgFyOROoFo="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"HeRi25TEvRPHRhdlmEwAVqP5vrSzcLkSVV70rCoeGYA="}, "signature":"MEUCIQCTAqna1zy/RXjepdrc7tjKVbgz962uTr7Gs0LxvF30IgIgKSHQG0f6dBiXsT0IkZX5H4wlu+ZvAlTAuFdliCHvbAU="}} \ No newline at end of file diff --git a/build/torch212-cxx11-cu132-x86_64-linux/modules.py b/build/torch212-cxx11-cu132-x86_64-linux/modules.py new file mode 100644 index 0000000000000000000000000000000000000000..222e384f2151c6a32275abcdb5000d5b8f1379d4 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/modules.py @@ -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, + ) diff --git a/build/torch212-cxx11-cu132-x86_64-linux/natten/__init__.py b/build/torch212-cxx11-cu132-x86_64-linux/natten/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/natten/__init__.py @@ -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"))) diff --git a/build/torch212-cxx11-cu132-x86_64-linux/token_permute/__init__.py b/build/torch212-cxx11-cu132-x86_64-linux/token_permute/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..381de37e4e7d4d2e6158668fbcc2ca52f7b388a3 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/token_permute/__init__.py @@ -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", +] diff --git a/build/torch212-cxx11-cu132-x86_64-linux/token_permute/cutlass_impl.py b/build/torch212-cxx11-cu132-x86_64-linux/token_permute/cutlass_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..49ce678f9052679448f843922f16bdec5d8796be --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/token_permute/cutlass_impl.py @@ -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 diff --git a/build/torch212-cxx11-cu132-x86_64-linux/token_permute/frontend.py b/build/torch212-cxx11-cu132-x86_64-linux/token_permute/frontend.py new file mode 100644 index 0000000000000000000000000000000000000000..720aded02e93e31f19a91efbda82cac3db15e4d2 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/token_permute/frontend.py @@ -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 diff --git a/build/torch212-cxx11-cu132-x86_64-linux/token_permute/torch_impl.py b/build/torch212-cxx11-cu132-x86_64-linux/token_permute/torch_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..7f06204e57cb6fc96ccb9e10ac9f798ac2d0e7a4 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/token_permute/torch_impl.py @@ -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", +] diff --git a/build/torch212-cxx11-cu132-x86_64-linux/utils/__init__.py b/build/torch212-cxx11-cu132-x86_64-linux/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..99f6c1e8cca79bfdf04640b8a92602b205a407a0 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/utils/__init__.py @@ -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. +# +################################################################################################# diff --git a/build/torch212-cxx11-cu132-x86_64-linux/utils/checks.py b/build/torch212-cxx11-cu132-x86_64-linux/utils/checks.py new file mode 100644 index 0000000000000000000000000000000000000000..f58db50a581e0bc6b43e8de0fbccd9404beaa1dc --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/utils/checks.py @@ -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, + ) diff --git a/build/torch212-cxx11-cu132-x86_64-linux/utils/device.py b/build/torch212-cxx11-cu132-x86_64-linux/utils/device.py new file mode 100644 index 0000000000000000000000000000000000000000..7997ab474197c78f3d41138f17ae311d6e06b1cf --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/utils/device.py @@ -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 diff --git a/build/torch212-cxx11-cu132-x86_64-linux/utils/dtype.py b/build/torch212-cxx11-cu132-x86_64-linux/utils/dtype.py new file mode 100644 index 0000000000000000000000000000000000000000..c65d9e9e0082c6f14af0b8331216cf26fcb3f28d --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/utils/dtype.py @@ -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] diff --git a/build/torch212-cxx11-cu132-x86_64-linux/utils/environment.py b/build/torch212-cxx11-cu132-x86_64-linux/utils/environment.py new file mode 100644 index 0000000000000000000000000000000000000000..0581aacb1183a5dea228a63859837113429e1a62 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/utils/environment.py @@ -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 diff --git a/build/torch212-cxx11-cu132-x86_64-linux/utils/log.py b/build/torch212-cxx11-cu132-x86_64-linux/utils/log.py new file mode 100644 index 0000000000000000000000000000000000000000..138c3ac8763480c0a381a8af0d913efcfa48b1a1 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/utils/log.py @@ -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) diff --git a/build/torch212-cxx11-cu132-x86_64-linux/utils/tensor.py b/build/torch212-cxx11-cu132-x86_64-linux/utils/tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..32d95acf2eb9a8c6fa29c7d3de3b48b5f14be415 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/utils/tensor.py @@ -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 diff --git a/build/torch212-cxx11-cu132-x86_64-linux/utils/testing.py b/build/torch212-cxx11-cu132-x86_64-linux/utils/testing.py new file mode 100644 index 0000000000000000000000000000000000000000..8ed9de5df24754dc9207a0b39954c8a73d7610a5 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/utils/testing.py @@ -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 diff --git a/build/torch212-cxx11-cu132-x86_64-linux/utils/tuples.py b/build/torch212-cxx11-cu132-x86_64-linux/utils/tuples.py new file mode 100644 index 0000000000000000000000000000000000000000..abf8a862fa4b647aab55fb2982e5dea08609d316 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/utils/tuples.py @@ -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 diff --git a/build/torch212-cxx11-cu132-x86_64-linux/utils/varlen.py b/build/torch212-cxx11-cu132-x86_64-linux/utils/varlen.py new file mode 100644 index 0000000000000000000000000000000000000000..8bc541dbd8e80ddf9923a097ac473b31035174ef --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/utils/varlen.py @@ -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, + ) diff --git a/build/torch212-cxx11-cu132-x86_64-linux/version.py b/build/torch212-cxx11-cu132-x86_64-linux/version.py new file mode 100644 index 0000000000000000000000000000000000000000..fe22ed72b87e6b242b01aae1ffb3de069ae7b0a9 --- /dev/null +++ b/build/torch212-cxx11-cu132-x86_64-linux/version.py @@ -0,0 +1,24 @@ +################################################################################################# +# 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. +# +################################################################################################# + +__version__ = "0.21.7" diff --git a/build/torch213-cxx11-cu126-x86_64-linux/__init__.py b/build/torch213-cxx11-cu126-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9fe10d400adaa8333f0da9c0c9654864659fbc3b --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/__init__.py @@ -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", +] diff --git a/build/torch213-cxx11-cu126-x86_64-linux/_environment.py b/build/torch213-cxx11-cu126-x86_64-linux/_environment.py new file mode 100644 index 0000000000000000000000000000000000000000..caae85bbe3c5eef992e4c6758e44b382382ebe41 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/_environment.py @@ -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", +] diff --git a/build/torch213-cxx11-cu126-x86_64-linux/_libnatten/__init__.py b/build/torch213-cxx11-cu126-x86_64-linux/_libnatten/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9566f8d4e7d9356415bb2951613dc5ada194dbd9 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/_libnatten/__init__.py @@ -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", +] diff --git a/build/torch213-cxx11-cu126-x86_64-linux/_libnatten/torch_wrappers.py b/build/torch213-cxx11-cu126-x86_64-linux/_libnatten/torch_wrappers.py new file mode 100644 index 0000000000000000000000000000000000000000..16e5b62b2bc5de0810f58d82d70d7fb2e3e56e4a --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/_libnatten/torch_wrappers.py @@ -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", +] diff --git a/build/torch213-cxx11-cu126-x86_64-linux/_natten_cuda_3641131.abi3.so b/build/torch213-cxx11-cu126-x86_64-linux/_natten_cuda_3641131.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..0f2f3d897992a02c550195a94964c0ea680ce40b --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/_natten_cuda_3641131.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:dc82de9264fed7caeae4e64b7b9aa7ada08ef5b11f8dff2df205ac89101264ea +size 103602504 diff --git a/build/torch213-cxx11-cu126-x86_64-linux/_ops.py b/build/torch213-cxx11-cu126-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..ee2d34306175408457b5e3dca950160c1444f947 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/_ops.py @@ -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}" diff --git a/build/torch213-cxx11-cu126-x86_64-linux/_types.py b/build/torch213-cxx11-cu126-x86_64-linux/_types.py new file mode 100644 index 0000000000000000000000000000000000000000..804980900b74586f5861f473853055c7d9630efd --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/_types.py @@ -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 diff --git a/build/torch213-cxx11-cu126-x86_64-linux/attn_merge.py b/build/torch213-cxx11-cu126-x86_64-linux/attn_merge.py new file mode 100644 index 0000000000000000000000000000000000000000..e71cc799d407df0afa701942ad08b03718ba20bf --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/attn_merge.py @@ -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"] diff --git a/build/torch213-cxx11-cu126-x86_64-linux/backends/__init__.py b/build/torch213-cxx11-cu126-x86_64-linux/backends/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..84b8192a37e048880aba2059a6f47e7d26f6a8ee --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/backends/__init__.py @@ -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", +] diff --git a/build/torch213-cxx11-cu126-x86_64-linux/backends/blackwell_fmha.py b/build/torch213-cxx11-cu126-x86_64-linux/backends/blackwell_fmha.py new file mode 100644 index 0000000000000000000000000000000000000000..9b36e79d8b9242b7c5a20833932721556abd61d5 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/backends/blackwell_fmha.py @@ -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 diff --git a/build/torch213-cxx11-cu126-x86_64-linux/backends/blackwell_fna.py b/build/torch213-cxx11-cu126-x86_64-linux/backends/blackwell_fna.py new file mode 100644 index 0000000000000000000000000000000000000000..5bd6f0ed8dd3da85bd60804e503f2aca27b37e4e --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/backends/blackwell_fna.py @@ -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, + ) diff --git a/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/__init__.py b/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9c2b75f7e5007dd1490e8ff5d4c6413f21a5e39f --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/__init__.py @@ -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) diff --git a/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/checks.py b/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/checks.py new file mode 100644 index 0000000000000000000000000000000000000000..4237d8eba4f2e11f7f7abfc796d53b87801f34a6 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/checks.py @@ -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 diff --git a/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass/__init__.py b/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fdd2d4e18d0d959b16f7c9bd0a8b2f71081845be --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass/__init__.py @@ -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}" + ) diff --git a/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass/backward_knobs.py b/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass/backward_knobs.py new file mode 100644 index 0000000000000000000000000000000000000000..dfb1c10a3f1dfa28597f6eb30e5f925ad4785325 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass/backward_knobs.py @@ -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.") diff --git a/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_backward_128x128.py b/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_backward_128x128.py new file mode 100644 index 0000000000000000000000000000000000000000..af6047ccb59f6923be82d6379b92c39b9bcb7f1a --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_backward_128x128.py @@ -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)), + ], +} diff --git a/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_backward_128x64.py b/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_backward_128x64.py new file mode 100644 index 0000000000000000000000000000000000000000..6a915e9686203919079092eccc743dce88b4107a --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_backward_128x64.py @@ -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)), + ], +} diff --git a/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_backward_64x64.py b/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_backward_64x64.py new file mode 100644 index 0000000000000000000000000000000000000000..8511bd844f34a263d233d4e6a1f6f7c775139e2d --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_backward_64x64.py @@ -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)), + ], +} diff --git a/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_forward_32x128.py b/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_forward_32x128.py new file mode 100644 index 0000000000000000000000000000000000000000..a930aba67df1de878d7e95797bb3b6d2d903e020 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_forward_32x128.py @@ -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)), + ], +} diff --git a/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_forward_64x128.py b/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_forward_64x128.py new file mode 100644 index 0000000000000000000000000000000000000000..1bc1fc78e028a755b755ea8108264eccf0bb9659 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_forward_64x128.py @@ -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)), + ], +} diff --git a/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_forward_64x64.py b/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_forward_64x64.py new file mode 100644 index 0000000000000000000000000000000000000000..c321c5b4e1a96f0cb58894f95dd2826d2d21c950 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass/fna_forward_64x64.py @@ -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)), + ], +} diff --git a/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass_blackwell/__init__.py b/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass_blackwell/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4a8c87c6c6d916f58c533fc2a84c66aff8bfe46b --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass_blackwell/__init__.py @@ -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}" + ) diff --git a/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass_hopper/__init__.py b/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass_hopper/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..70b2e3106003ed583050be356b54e32c9b7c61c0 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/cutlass_hopper/__init__.py @@ -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}" + ) diff --git a/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/flex/__init__.py b/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/flex/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c5027701152d3d5976670fd59dc158331c5cadfa --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/backends/configs/flex/__init__.py @@ -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}" + ) diff --git a/build/torch213-cxx11-cu126-x86_64-linux/backends/flex.py b/build/torch213-cxx11-cu126-x86_64-linux/backends/flex.py new file mode 100644 index 0000000000000000000000000000000000000000..0555a7ebdb62e6e4060fbfe3106bed116b44797b --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/backends/flex.py @@ -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, + ) diff --git a/build/torch213-cxx11-cu126-x86_64-linux/backends/fmha.py b/build/torch213-cxx11-cu126-x86_64-linux/backends/fmha.py new file mode 100644 index 0000000000000000000000000000000000000000..0c6851f8bb362a4c7d45db46e61db1ed1dc61d51 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/backends/fmha.py @@ -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 diff --git a/build/torch213-cxx11-cu126-x86_64-linux/backends/fna.py b/build/torch213-cxx11-cu126-x86_64-linux/backends/fna.py new file mode 100644 index 0000000000000000000000000000000000000000..747ecb1f7dcfe032e02e4fa76b8a3a9b70e96bd0 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/backends/fna.py @@ -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, + ) diff --git a/build/torch213-cxx11-cu126-x86_64-linux/backends/hopper_fmha.py b/build/torch213-cxx11-cu126-x86_64-linux/backends/hopper_fmha.py new file mode 100644 index 0000000000000000000000000000000000000000..251b50dc34598111f288ed3ac3d2b414ea23f253 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/backends/hopper_fmha.py @@ -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 diff --git a/build/torch213-cxx11-cu126-x86_64-linux/backends/hopper_fna.py b/build/torch213-cxx11-cu126-x86_64-linux/backends/hopper_fna.py new file mode 100644 index 0000000000000000000000000000000000000000..17bb6d5952675043b66b59fcf92ea9eea54facf9 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/backends/hopper_fna.py @@ -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, + ) diff --git a/build/torch213-cxx11-cu126-x86_64-linux/backends/reference.py b/build/torch213-cxx11-cu126-x86_64-linux/backends/reference.py new file mode 100644 index 0000000000000000000000000000000000000000..d0f34b41a8c9b1a1051360b23d7e20573931bb27 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/backends/reference.py @@ -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, + ) diff --git a/build/torch213-cxx11-cu126-x86_64-linux/context.py b/build/torch213-cxx11-cu126-x86_64-linux/context.py new file mode 100644 index 0000000000000000000000000000000000000000..1793ac984bbe5be383bec2c07b54c770c9969b42 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/context.py @@ -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) diff --git a/build/torch213-cxx11-cu126-x86_64-linux/functional.py b/build/torch213-cxx11-cu126-x86_64-linux/functional.py new file mode 100644 index 0000000000000000000000000000000000000000..944a0c2a7786de250fd5031e933133ea676d2a79 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/functional.py @@ -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, + ) diff --git a/build/torch213-cxx11-cu126-x86_64-linux/metadata.json b/build/torch213-cxx11-cu126-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..1b1d2b8c4a9bf8e95fd300914f35f49ec1f797c9 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/metadata.json @@ -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": "3ILekmT+18rq5OZLe5qnraCO9bEfjf8t8gWsiRASZOo=", + "_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 + } + } +} \ No newline at end of file diff --git a/build/torch213-cxx11-cu126-x86_64-linux/metadata.json.sigstore b/build/torch213-cxx11-cu126-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..280b6cc3024f49b18dae2ed0339bec7e4722812c --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSTCCBtCgAwIBAgIUIJ0eMYy3MT5sFppHZlmMnbTxX78wCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNzI5MDkyMTM1WhcNMjYwNzI5MDkzMTM1WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEksIrMElvogpp+nuqfIjjlVNLEcDMAt5Nx3gblDwys3h3+4AtNxvSqE0wCrDr9ifDw0YOaqsR68Z86eZKHEXPPqOCBe8wggXrMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUnMkc1jfz5qUNWKmfPWhw0pWnH7IwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoMjBhNjEzYzEwZDliMWM5Y2I0M2E1MWU4NTBlOWZhMDQ2OWE5NzE5OTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoMjBhNjEzYzEwZDliMWM5Y2I0M2E1MWU4NTBlOWZhMDQ2OWE5NzE5OTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoMjBhNjEzYzEwZDliMWM5Y2I0M2E1MWU4NTBlOWZhMDQ2OWE5NzE5OTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDIwYTYxM2MxMGQ5YjFjOWNiNDNhNTFlODUwZTlmYTA0NjlhOTcxOTkwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzA0MTgyNTI4NjkvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiQYKKwYBBAHWeQIEAgR7BHkAdwB1AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABn60t7LQAAAQDAEYwRAIgN8pFtsSg0PmxreAgfDVN3uZdY6EWGvKx7pTicTxgWf8CIFWnA0rzfVlZarJxJTRDLUbdYxg5BFvcMgt6ppVKxsDrMAoGCCqGSM49BAMDA2cAMGQCMFYmL2OrbreFsfuhLJeopWbl2GJ12ryWUZ0Viy9+qO5Z71OrkpNy2aI6qXQDFgKbHQIwfan0DYNXTlvMtZifGAjwU5ZN9D+WLd9bD66PO7Vdc8kdHlMd6oLuKd009d1rlMbX"}, "tlogEntries":[{"logIndex":"2280149253", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1785316896", "inclusionPromise":{"signedEntryTimestamp":"MEUCIQCJ8kAJjALw9EEDJuqnHdeaFF3KfmxdNK42+38A4t5ZrgIgcssBXifN7MEmUhloAiu7yjyFF652GkKG5ujHepoYflg="}, "inclusionProof":{"logIndex":"2158244991", "rootHash":"ALkiW0TQf3vKt+a4df0Fuk8chwgqSMm4k3Jv450PJns=", "treeSize":"2158244992", "hashes":["Q236MpYFUe9g9PneRsMy7CVQWqe/19w/WsPGDjNvsng=", "It9ePlHtkB2u9rVMQvt8DmLeLnu5oGkAsh/Vo9Z0B6s=", "eU6Cx6yRPX1KVMKCuoTLtpTAJgcu8gQ5PEaj5SxFle8=", "kiHFfzHEpmkWpNyU+EGq+/Uh1mX899tnYY9rhav4eNY=", "3LrZZ5T93/7CInZzUe+S+IcEcc7ONboRD4Efht9iG/Y=", "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\n2158244992\nALkiW0TQf3vKt+a4df0Fuk8chwgqSMm4k3Jv450PJns=\n\n— rekor.sigstore.dev wNI9ajBFAiA/fry2xzOhphY3tZ+79wn3hX+v2XgvI16RV5NEcNi4KAIhANgNK2Th3a5XznGn6CvNGYh3xQLw7N9cO4oAL2C5soRE\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiIwNDY1OTRjMmIyNjAzY2FjODNkOWE3ZmU1NDEzOGI0YTc0MTdiM2ZmMWMzNmFlOGUwZDQ2YWFmNTc4ODIxMjhiIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUNjSkVrVVFNOFcwRXNRQVFEUWhvSkw3bUFuQmJlcWlBVEJCNm8zanVRUEhnSWdPaFpVMzJLWFp2RGh5enFkNEcwMUFmUXdtWlRiZWoxaGtTOFA5TCtMTjRzPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRWRU5EUW5SRFowRjNTVUpCWjBsVlNVb3daVTFaZVROTlZEVnpSbkJ3U0Zwc2JVMXVZbFI0V0RjNGQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDU2U1RWTlJHdDVUVlJOTVZkb1kwNU5hbGwzVG5wSk5VMUVhM3BOVkUweFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZyYzBseVRVVnNkbTluY0hBcmJuVnhaa2xxYW14V1RreEZZMFJOUVhRMVRuZ3paMklLYkVSM2VYTXphRE1yTkVGMFRuaDJVM0ZGTUhkRGNrUnlPV2xtUkhjd1dVOWhjWE5TTmpoYU9EWmxXa3RJUlZoUVVIRlBRMEpsT0hkbloxaHlUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZ1VFd0akNqRnFabm8xY1ZWT1YwdHRabEJYYUhjd2NGZHVTRGRKZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU5ha0pvVG1wRmVsbDZSWGRhUkd4cFRWZE5OVmt5U1RCTk1rVXhDazFYVlRST1ZFSnNUMWRhYUUxRVVUSlBWMFUxVG5wRk5VOVVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDFxUW1oT2FrVjZXWHBGZDFwRWJHbE5WMDAxV1RKSk1FMHlSVEZOVjFVMFRsUkNiRTlYV21nS1RVUlJNazlYUlRWT2VrVTFUMVJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDFxUW1oT2FrVjZXWHBGZDFwRWJHbE5WMDAxV1RKSk1FMHlSVEVLVFZkVk5FNVVRbXhQVjFwb1RVUlJNazlYUlRWT2VrVTFUMVJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUkpkMWxVV1hnS1RUSk5lRTFIVVRWWmFrWnFUMWRPYVU1RVRtaE9WRVpzVDBSVmQxcFViRzFaVkVFd1RtcHNhRTlVWTNoUFZHdDNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWtFd1RWUm5lVTVVU1RST2FtdDJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwVVZsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTTjBKSWEwRUtaSGRDTVVGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNDJNSFEzVEZGQlFVRlJSQXBCUlZsM1VrRkpaMDQ0Y0VaMGMxTm5NRkJ0ZUhKbFFXZG1SRlpPTTNWYVpGazJSVmRIZGt0NE4zQlVhV05VZUdkWFpqaERTVVpYYmtFd2NucG1WbXhhQ21GeVNuaEtWRkpFVEZWaVpGbDRaelZDUm5aalRXZDBObkJ3Vmt0NGMwUnlUVUZ2UjBORGNVZFRUVFE1UWtGTlJFRXlZMEZOUjFGRFRVWlpiVXd5VDNJS1luSmxSbk5tZFdoTVNtVnZjRmRpYkRKSFNqRXljbmxYVlZvd1ZtbDVPU3R4VHpWYU56RlBjbXR3VG5reVlVazJjVmhSUkVablMySklVVWwzWm1GdU1BcEVXVTVZVkd4MlRYUmFhV1pIUVdwM1ZUVmFUamxFSzFkTVpEbGlSRFkyVUU4M1ZtUmpPR3RrU0d4TlpEWnZUSFZMWkRBd09XUXhjbXhOWWxnS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgvgaVgSph+O/aqDu3E/j4LU5hEkqLlZtczcdLSAeDJe4CFQC0+K1/+l6cE48f9ncTj9zXd5V1LRgPMjAyNjA3MjkwOTIxMzZaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHbMIIB1wIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDcyOTA5MjEzNlowLwYJKoZIhvcNAQkEMSIEILo/4KTMof+HejMUHkNG7JAK0hrAjtc5JiRuoe9pWZFzMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRnMGUCMQDPdx156QmlIIUVSOnTN578VzuvIq43TbTzfMenoTIvrLdy9Zaj2vtdQY1KFXrE0IICMBDVKynrZHqa4gpaGYhEQmDcX1EOMCWkjfFt5Sq1HzJH14RyoSUdmJfGUFPY8VhbDQ=="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"BGWUwrJgPKyD2af+VBOLSnQXs/8cNq6ODUaq9XiCEos="}, "signature":"MEUCIQCcJEkUQM8W0EsQAQDQhoJL7mAnBbeqiATBB6o3juQPHgIgOhZU32KXZvDhyzqd4G01AfQwmZTbej1hkS8P9L+LN4s="}} \ No newline at end of file diff --git a/build/torch213-cxx11-cu126-x86_64-linux/modules.py b/build/torch213-cxx11-cu126-x86_64-linux/modules.py new file mode 100644 index 0000000000000000000000000000000000000000..222e384f2151c6a32275abcdb5000d5b8f1379d4 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/modules.py @@ -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, + ) diff --git a/build/torch213-cxx11-cu126-x86_64-linux/natten/__init__.py b/build/torch213-cxx11-cu126-x86_64-linux/natten/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/natten/__init__.py @@ -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"))) diff --git a/build/torch213-cxx11-cu126-x86_64-linux/token_permute/__init__.py b/build/torch213-cxx11-cu126-x86_64-linux/token_permute/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..381de37e4e7d4d2e6158668fbcc2ca52f7b388a3 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/token_permute/__init__.py @@ -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", +] diff --git a/build/torch213-cxx11-cu126-x86_64-linux/token_permute/cutlass_impl.py b/build/torch213-cxx11-cu126-x86_64-linux/token_permute/cutlass_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..49ce678f9052679448f843922f16bdec5d8796be --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/token_permute/cutlass_impl.py @@ -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 diff --git a/build/torch213-cxx11-cu126-x86_64-linux/token_permute/frontend.py b/build/torch213-cxx11-cu126-x86_64-linux/token_permute/frontend.py new file mode 100644 index 0000000000000000000000000000000000000000..720aded02e93e31f19a91efbda82cac3db15e4d2 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/token_permute/frontend.py @@ -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 diff --git a/build/torch213-cxx11-cu126-x86_64-linux/token_permute/torch_impl.py b/build/torch213-cxx11-cu126-x86_64-linux/token_permute/torch_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..7f06204e57cb6fc96ccb9e10ac9f798ac2d0e7a4 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/token_permute/torch_impl.py @@ -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", +] diff --git a/build/torch213-cxx11-cu126-x86_64-linux/utils/__init__.py b/build/torch213-cxx11-cu126-x86_64-linux/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..99f6c1e8cca79bfdf04640b8a92602b205a407a0 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/utils/__init__.py @@ -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. +# +################################################################################################# diff --git a/build/torch213-cxx11-cu126-x86_64-linux/utils/checks.py b/build/torch213-cxx11-cu126-x86_64-linux/utils/checks.py new file mode 100644 index 0000000000000000000000000000000000000000..f58db50a581e0bc6b43e8de0fbccd9404beaa1dc --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/utils/checks.py @@ -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, + ) diff --git a/build/torch213-cxx11-cu126-x86_64-linux/utils/device.py b/build/torch213-cxx11-cu126-x86_64-linux/utils/device.py new file mode 100644 index 0000000000000000000000000000000000000000..7997ab474197c78f3d41138f17ae311d6e06b1cf --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/utils/device.py @@ -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 diff --git a/build/torch213-cxx11-cu126-x86_64-linux/utils/dtype.py b/build/torch213-cxx11-cu126-x86_64-linux/utils/dtype.py new file mode 100644 index 0000000000000000000000000000000000000000..c65d9e9e0082c6f14af0b8331216cf26fcb3f28d --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/utils/dtype.py @@ -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] diff --git a/build/torch213-cxx11-cu126-x86_64-linux/utils/environment.py b/build/torch213-cxx11-cu126-x86_64-linux/utils/environment.py new file mode 100644 index 0000000000000000000000000000000000000000..0581aacb1183a5dea228a63859837113429e1a62 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/utils/environment.py @@ -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 diff --git a/build/torch213-cxx11-cu126-x86_64-linux/utils/log.py b/build/torch213-cxx11-cu126-x86_64-linux/utils/log.py new file mode 100644 index 0000000000000000000000000000000000000000..138c3ac8763480c0a381a8af0d913efcfa48b1a1 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/utils/log.py @@ -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) diff --git a/build/torch213-cxx11-cu126-x86_64-linux/utils/tensor.py b/build/torch213-cxx11-cu126-x86_64-linux/utils/tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..32d95acf2eb9a8c6fa29c7d3de3b48b5f14be415 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/utils/tensor.py @@ -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 diff --git a/build/torch213-cxx11-cu126-x86_64-linux/utils/testing.py b/build/torch213-cxx11-cu126-x86_64-linux/utils/testing.py new file mode 100644 index 0000000000000000000000000000000000000000..8ed9de5df24754dc9207a0b39954c8a73d7610a5 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/utils/testing.py @@ -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 diff --git a/build/torch213-cxx11-cu126-x86_64-linux/utils/tuples.py b/build/torch213-cxx11-cu126-x86_64-linux/utils/tuples.py new file mode 100644 index 0000000000000000000000000000000000000000..abf8a862fa4b647aab55fb2982e5dea08609d316 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/utils/tuples.py @@ -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 diff --git a/build/torch213-cxx11-cu126-x86_64-linux/utils/varlen.py b/build/torch213-cxx11-cu126-x86_64-linux/utils/varlen.py new file mode 100644 index 0000000000000000000000000000000000000000..8bc541dbd8e80ddf9923a097ac473b31035174ef --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/utils/varlen.py @@ -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, + ) diff --git a/build/torch213-cxx11-cu126-x86_64-linux/version.py b/build/torch213-cxx11-cu126-x86_64-linux/version.py new file mode 100644 index 0000000000000000000000000000000000000000..fe22ed72b87e6b242b01aae1ffb3de069ae7b0a9 --- /dev/null +++ b/build/torch213-cxx11-cu126-x86_64-linux/version.py @@ -0,0 +1,24 @@ +################################################################################################# +# 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. +# +################################################################################################# + +__version__ = "0.21.7" diff --git a/build/torch213-cxx11-cu130-x86_64-linux/__init__.py b/build/torch213-cxx11-cu130-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9fe10d400adaa8333f0da9c0c9654864659fbc3b --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/__init__.py @@ -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", +] diff --git a/build/torch213-cxx11-cu130-x86_64-linux/_environment.py b/build/torch213-cxx11-cu130-x86_64-linux/_environment.py new file mode 100644 index 0000000000000000000000000000000000000000..caae85bbe3c5eef992e4c6758e44b382382ebe41 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/_environment.py @@ -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", +] diff --git a/build/torch213-cxx11-cu130-x86_64-linux/_libnatten/__init__.py b/build/torch213-cxx11-cu130-x86_64-linux/_libnatten/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9566f8d4e7d9356415bb2951613dc5ada194dbd9 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/_libnatten/__init__.py @@ -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", +] diff --git a/build/torch213-cxx11-cu130-x86_64-linux/_libnatten/torch_wrappers.py b/build/torch213-cxx11-cu130-x86_64-linux/_libnatten/torch_wrappers.py new file mode 100644 index 0000000000000000000000000000000000000000..16e5b62b2bc5de0810f58d82d70d7fb2e3e56e4a --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/_libnatten/torch_wrappers.py @@ -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", +] diff --git a/build/torch213-cxx11-cu130-x86_64-linux/_natten_cuda_3641131.abi3.so b/build/torch213-cxx11-cu130-x86_64-linux/_natten_cuda_3641131.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..87701d8e42fe114864b1bcb5b2aa56c9e698953a --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/_natten_cuda_3641131.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:98f996db5d78ff428992c8288bb13a1b0ca14fba332c2900fec599d436c39d3f +size 164303768 diff --git a/build/torch213-cxx11-cu130-x86_64-linux/_ops.py b/build/torch213-cxx11-cu130-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..ee2d34306175408457b5e3dca950160c1444f947 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/_ops.py @@ -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}" diff --git a/build/torch213-cxx11-cu130-x86_64-linux/_types.py b/build/torch213-cxx11-cu130-x86_64-linux/_types.py new file mode 100644 index 0000000000000000000000000000000000000000..804980900b74586f5861f473853055c7d9630efd --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/_types.py @@ -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 diff --git a/build/torch213-cxx11-cu130-x86_64-linux/attn_merge.py b/build/torch213-cxx11-cu130-x86_64-linux/attn_merge.py new file mode 100644 index 0000000000000000000000000000000000000000..e71cc799d407df0afa701942ad08b03718ba20bf --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/attn_merge.py @@ -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"] diff --git a/build/torch213-cxx11-cu130-x86_64-linux/backends/__init__.py b/build/torch213-cxx11-cu130-x86_64-linux/backends/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..84b8192a37e048880aba2059a6f47e7d26f6a8ee --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/backends/__init__.py @@ -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", +] diff --git a/build/torch213-cxx11-cu130-x86_64-linux/backends/blackwell_fmha.py b/build/torch213-cxx11-cu130-x86_64-linux/backends/blackwell_fmha.py new file mode 100644 index 0000000000000000000000000000000000000000..9b36e79d8b9242b7c5a20833932721556abd61d5 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/backends/blackwell_fmha.py @@ -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 diff --git a/build/torch213-cxx11-cu130-x86_64-linux/backends/blackwell_fna.py b/build/torch213-cxx11-cu130-x86_64-linux/backends/blackwell_fna.py new file mode 100644 index 0000000000000000000000000000000000000000..5bd6f0ed8dd3da85bd60804e503f2aca27b37e4e --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/backends/blackwell_fna.py @@ -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, + ) diff --git a/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/__init__.py b/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9c2b75f7e5007dd1490e8ff5d4c6413f21a5e39f --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/__init__.py @@ -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) diff --git a/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/checks.py b/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/checks.py new file mode 100644 index 0000000000000000000000000000000000000000..4237d8eba4f2e11f7f7abfc796d53b87801f34a6 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/checks.py @@ -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 diff --git a/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass/__init__.py b/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fdd2d4e18d0d959b16f7c9bd0a8b2f71081845be --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass/__init__.py @@ -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}" + ) diff --git a/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass/backward_knobs.py b/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass/backward_knobs.py new file mode 100644 index 0000000000000000000000000000000000000000..dfb1c10a3f1dfa28597f6eb30e5f925ad4785325 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass/backward_knobs.py @@ -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.") diff --git a/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_backward_128x128.py b/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_backward_128x128.py new file mode 100644 index 0000000000000000000000000000000000000000..af6047ccb59f6923be82d6379b92c39b9bcb7f1a --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_backward_128x128.py @@ -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)), + ], +} diff --git a/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_backward_128x64.py b/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_backward_128x64.py new file mode 100644 index 0000000000000000000000000000000000000000..6a915e9686203919079092eccc743dce88b4107a --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_backward_128x64.py @@ -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)), + ], +} diff --git a/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_backward_64x64.py b/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_backward_64x64.py new file mode 100644 index 0000000000000000000000000000000000000000..8511bd844f34a263d233d4e6a1f6f7c775139e2d --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_backward_64x64.py @@ -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)), + ], +} diff --git a/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_forward_32x128.py b/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_forward_32x128.py new file mode 100644 index 0000000000000000000000000000000000000000..a930aba67df1de878d7e95797bb3b6d2d903e020 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_forward_32x128.py @@ -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)), + ], +} diff --git a/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_forward_64x128.py b/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_forward_64x128.py new file mode 100644 index 0000000000000000000000000000000000000000..1bc1fc78e028a755b755ea8108264eccf0bb9659 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_forward_64x128.py @@ -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)), + ], +} diff --git a/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_forward_64x64.py b/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_forward_64x64.py new file mode 100644 index 0000000000000000000000000000000000000000..c321c5b4e1a96f0cb58894f95dd2826d2d21c950 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass/fna_forward_64x64.py @@ -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)), + ], +} diff --git a/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass_blackwell/__init__.py b/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass_blackwell/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4a8c87c6c6d916f58c533fc2a84c66aff8bfe46b --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass_blackwell/__init__.py @@ -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}" + ) diff --git a/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass_hopper/__init__.py b/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass_hopper/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..70b2e3106003ed583050be356b54e32c9b7c61c0 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/cutlass_hopper/__init__.py @@ -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}" + ) diff --git a/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/flex/__init__.py b/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/flex/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c5027701152d3d5976670fd59dc158331c5cadfa --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/backends/configs/flex/__init__.py @@ -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}" + ) diff --git a/build/torch213-cxx11-cu130-x86_64-linux/backends/flex.py b/build/torch213-cxx11-cu130-x86_64-linux/backends/flex.py new file mode 100644 index 0000000000000000000000000000000000000000..0555a7ebdb62e6e4060fbfe3106bed116b44797b --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/backends/flex.py @@ -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, + ) diff --git a/build/torch213-cxx11-cu130-x86_64-linux/backends/fmha.py b/build/torch213-cxx11-cu130-x86_64-linux/backends/fmha.py new file mode 100644 index 0000000000000000000000000000000000000000..0c6851f8bb362a4c7d45db46e61db1ed1dc61d51 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/backends/fmha.py @@ -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 diff --git a/build/torch213-cxx11-cu130-x86_64-linux/backends/fna.py b/build/torch213-cxx11-cu130-x86_64-linux/backends/fna.py new file mode 100644 index 0000000000000000000000000000000000000000..747ecb1f7dcfe032e02e4fa76b8a3a9b70e96bd0 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/backends/fna.py @@ -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, + ) diff --git a/build/torch213-cxx11-cu130-x86_64-linux/backends/hopper_fmha.py b/build/torch213-cxx11-cu130-x86_64-linux/backends/hopper_fmha.py new file mode 100644 index 0000000000000000000000000000000000000000..251b50dc34598111f288ed3ac3d2b414ea23f253 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/backends/hopper_fmha.py @@ -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 diff --git a/build/torch213-cxx11-cu130-x86_64-linux/backends/hopper_fna.py b/build/torch213-cxx11-cu130-x86_64-linux/backends/hopper_fna.py new file mode 100644 index 0000000000000000000000000000000000000000..17bb6d5952675043b66b59fcf92ea9eea54facf9 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/backends/hopper_fna.py @@ -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, + ) diff --git a/build/torch213-cxx11-cu130-x86_64-linux/backends/reference.py b/build/torch213-cxx11-cu130-x86_64-linux/backends/reference.py new file mode 100644 index 0000000000000000000000000000000000000000..d0f34b41a8c9b1a1051360b23d7e20573931bb27 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/backends/reference.py @@ -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, + ) diff --git a/build/torch213-cxx11-cu130-x86_64-linux/context.py b/build/torch213-cxx11-cu130-x86_64-linux/context.py new file mode 100644 index 0000000000000000000000000000000000000000..1793ac984bbe5be383bec2c07b54c770c9969b42 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/context.py @@ -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) diff --git a/build/torch213-cxx11-cu130-x86_64-linux/functional.py b/build/torch213-cxx11-cu130-x86_64-linux/functional.py new file mode 100644 index 0000000000000000000000000000000000000000..944a0c2a7786de250fd5031e933133ea676d2a79 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/functional.py @@ -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, + ) diff --git a/build/torch213-cxx11-cu130-x86_64-linux/metadata.json b/build/torch213-cxx11-cu130-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..d2d6f6b6fb371a4ca305a5568d202cb380ba89cd --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/metadata.json @@ -0,0 +1,84 @@ +{ + "name": "natten", + "id": "_natten_cuda_3641131", + "version": 1, + "license": "MIT", + "upstream": "https://github.com/SHI-Labs/NATTEN", + "python-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "10.0", + "10.0a", + "12.0", + "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": "mPmW2114/0KJksgoi7E6GwyhT7ozLCkA/sWZ1DbDnT8=", + "_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 + } + } +} \ No newline at end of file diff --git a/build/torch213-cxx11-cu130-x86_64-linux/metadata.json.sigstore b/build/torch213-cxx11-cu130-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..987e1df78e8b3385dd547bb30f616170be7ab0f6 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSjCCBtGgAwIBAgIUUIL8W6gKjbCf1rSfLX/TL8HnvM8wCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNzI5MDkyMTM4WhcNMjYwNzI5MDkzMTM4WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEwnORAXckNekRSe3j6sD/66vE2j6xhsjN+sC4/yaTSLJm22P1zMTsjubbvQbWDPHA4hexu2RnHqRZJk1YiXzkUqOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUgQ/mljuVgWmGiNbqAkfX2kE+9pgwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoMjBhNjEzYzEwZDliMWM5Y2I0M2E1MWU4NTBlOWZhMDQ2OWE5NzE5OTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoMjBhNjEzYzEwZDliMWM5Y2I0M2E1MWU4NTBlOWZhMDQ2OWE5NzE5OTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoMjBhNjEzYzEwZDliMWM5Y2I0M2E1MWU4NTBlOWZhMDQ2OWE5NzE5OTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDIwYTYxM2MxMGQ5YjFjOWNiNDNhNTFlODUwZTlmYTA0NjlhOTcxOTkwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzA0MTgyNTI4NjkvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABn60t9YMAAAQDAEcwRQIgGkxlhSd+pfTovUpr4caypnNMSQ3+SrMs2yH+nsjQyyACIQDNz17IW7FdSlx63zPs7TxmRojAc/hBVn7razQy6eKarzAKBggqhkjOPQQDAwNnADBkAjARE/JcZgfG+lFgY4iAqAMVxXRIHbvyCsRvf0Enc+z35B5AeyIlbGruvTcIFfAIBdoCMHTwHxBU+Zjwl0/6yS8FZ5RAgMQ4rvhge96P3JxfaUBH82jotNOfDfMbOY6e4xNW3A=="}, "tlogEntries":[{"logIndex":"2280149294", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1785316898", "inclusionPromise":{"signedEntryTimestamp":"MEUCIQCbfa6FXZA9QycGeUzcav0NMB6154VHTX0KfcgtFVMvKAIgKXsL/h0AyEuXFeA6LLR7iamzDScMp56JIn8gE92aLYI="}, "inclusionProof":{"logIndex":"2158245032", "rootHash":"nKUwg+zgNqjk0EYWelTn8GB53pi2CGsMfsyRVucNeKE=", "treeSize":"2158245033", "hashes":["EWp957f1+B2TjqrXnGYGbdw6PZkatBj0k8K6bB/5ZUY=", "p9fL2qSzXGY6RVJxNUNqpUili0XEEytLl9t5DXcnrqY=", "/5uu/jw8GTRTqEZ0auUV9lv+Zg0twg6hHe7ZChWeiJk=", "DwNueYXiyuVjob7c8CRNRN57ioovkJTOeOeOQIi6/pI=", "bXYdybsl/2BK7Yled3FQJEE3Hk0TeGcWnXUxOkVZv8w=", "xrwx0yHkH5ZmCgeWs/jLZZ8RcejqpeuUeAWNJkAYrHk=", "8MN6j2GHUjPnNtIwpas9l/XjdwE0/KPUhFDlrVkngko=", "b+xUZfuENQxvSOJxzNvYvRG8eVphfszPpZmuf4/cQ6c=", "OVsvZCKnWA+498QUIaQCtitUT6huDbC7SmhH1l8MxXI=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2158245033\nnKUwg+zgNqjk0EYWelTn8GB53pi2CGsMfsyRVucNeKE=\n\n— rekor.sigstore.dev wNI9ajBFAiAXiLSrW8Q+3DDGLOkSee8V/PY5MmVOJjnRwB5LLpjwYAIhAP/eUr6Dse4e6ZtUbGRMlnv2/Dq37kwhZiRXt8WA057m\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI3Yjg2MDkzYTJjNjNjZDQ0OTg2ZjMxYTVmMmM1MTU4NDRjZmVjYWVjMDEzZGVmNWNlM2IzYWU5M2JkNDcyM2ZlIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUURneWxxU3lhbnVQNGcyWGE0ZUFDVy9ubWhvWE8vOVVGbHh4MGF0b2dtcmFRSWdGSUlYZlU4c09IR0xzRzJOVDJRc3gxREVEd3l6WVpTNnVCdk5RMmt6Rm9jPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRha05EUW5SSFowRjNTVUpCWjBsVlZVbE1PRmMyWjB0cVlrTm1NWEpUWmt4WUwxUk1PRWh1ZGswNGQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDU2U1RWTlJHdDVUVlJOTkZkb1kwNU5hbGwzVG5wSk5VMUVhM3BOVkUwMFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVYzYms5U1FWaGphMDVsYTFKVFpUTnFObk5FTHpZMmRrVXlhalo0YUhOcVRpdHpRelFLTDNsaFZGTk1TbTB5TWxBeGVrMVVjMnAxWW1KMlVXSlhSRkJJUVRSb1pYaDFNbEp1U0hGU1drcHJNVmxwV0hwclZYRlBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZuVVM5dENteHFkVlpuVjIxSGFVNWljVUZyWmxneWEwVXJPWEJuZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU5ha0pvVG1wRmVsbDZSWGRhUkd4cFRWZE5OVmt5U1RCTk1rVXhDazFYVlRST1ZFSnNUMWRhYUUxRVVUSlBWMFUxVG5wRk5VOVVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDFxUW1oT2FrVjZXWHBGZDFwRWJHbE5WMDAxV1RKSk1FMHlSVEZOVjFVMFRsUkNiRTlYV21nS1RVUlJNazlYUlRWT2VrVTFUMVJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDFxUW1oT2FrVjZXWHBGZDFwRWJHbE5WMDAxV1RKSk1FMHlSVEVLVFZkVk5FNVVRbXhQVjFwb1RVUlJNazlYUlRWT2VrVTFUMVJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUkpkMWxVV1hnS1RUSk5lRTFIVVRWWmFrWnFUMWRPYVU1RVRtaE9WRVpzVDBSVmQxcFViRzFaVkVFd1RtcHNhRTlVWTNoUFZHdDNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWtFd1RWUm5lVTVVU1RST2FtdDJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNDJNSFE1V1UxQlFVRlJSQXBCUldOM1VsRkpaMGRyZUd4b1UyUXJjR1pVYjNaVmNISTBZMkY1Y0c1T1RWTlJNeXRUY2sxek1ubElLMjV6YWxGNWVVRkRTVkZFVG5veE4wbFhOMFprQ2xOc2VEWXplbEJ6TjFSNGJWSnZha0ZqTDJoQ1ZtNDNjbUY2VVhrMlpVdGhjbnBCUzBKblozRm9hMnBQVUZGUlJFRjNUbTVCUkVKclFXcEJVa1V2U21NS1dtZG1SeXRzUm1kWk5HbEJjVUZOVm5oWVVrbElZblo1UTNOU2RtWXdSVzVqSzNvek5VSTFRV1Y1U1d4aVIzSjFkbFJqU1VabVFVbENaRzlEVFVoVWR3cEllRUpWSzFwcWQyd3dMelo1VXpoR1dqVlNRV2ROVVRSeWRtaG5aVGsyVUROS2VHWmhWVUpJT0RKcWIzUk9UMlpFWmsxaVQxazJaVFI0VGxjelFUMDlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyDADAgEAMIICvwYJKoZIhvcNAQcCoIICsDCCAqwCAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgP6mcG0ztOoF1dpuczQrL3QeTdfE3Azlw73v1YTpPxzMCFAX6symJB2XBMNpBZRSJtTo40nWHGA8yMDI2MDcyOTA5MjEzOFowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdowggHWAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwNzI5MDkyMTM4WjAvBgkqhkiG9w0BCQQxIgQgkDvtCTAWp812s3k9cwgHLKOO7Mw4aj0XaYQm9/12S+QwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGYwZAIwVp/Sn63W9VsVV00MwtuxyAOvV1W7VBp+/LQ/kFbFMs0SIHNKVhzMBTGJi8Eyn3ifAjBxVhjAT09QDvvSFu5D+EnULAREVeJIUA/YeyAI7MUhIqNNyRtrvsGVXlrMAsNQ42o="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"e4YJOixjzUSYbzGl8sUVhEz+yuwBPe9c47Ouk71HI/4="}, "signature":"MEUCIQDgylqSyanuP4g2Xa4eACW/nmhoXO/9UFlxx0atogmraQIgFIIXfU8sOHGLsG2NT2Qsx1DEDwyzYZS6uBvNQ2kzFoc="}} \ No newline at end of file diff --git a/build/torch213-cxx11-cu130-x86_64-linux/modules.py b/build/torch213-cxx11-cu130-x86_64-linux/modules.py new file mode 100644 index 0000000000000000000000000000000000000000..222e384f2151c6a32275abcdb5000d5b8f1379d4 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/modules.py @@ -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, + ) diff --git a/build/torch213-cxx11-cu130-x86_64-linux/natten/__init__.py b/build/torch213-cxx11-cu130-x86_64-linux/natten/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/natten/__init__.py @@ -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"))) diff --git a/build/torch213-cxx11-cu130-x86_64-linux/token_permute/__init__.py b/build/torch213-cxx11-cu130-x86_64-linux/token_permute/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..381de37e4e7d4d2e6158668fbcc2ca52f7b388a3 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/token_permute/__init__.py @@ -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", +] diff --git a/build/torch213-cxx11-cu130-x86_64-linux/token_permute/cutlass_impl.py b/build/torch213-cxx11-cu130-x86_64-linux/token_permute/cutlass_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..49ce678f9052679448f843922f16bdec5d8796be --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/token_permute/cutlass_impl.py @@ -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 diff --git a/build/torch213-cxx11-cu130-x86_64-linux/token_permute/frontend.py b/build/torch213-cxx11-cu130-x86_64-linux/token_permute/frontend.py new file mode 100644 index 0000000000000000000000000000000000000000..720aded02e93e31f19a91efbda82cac3db15e4d2 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/token_permute/frontend.py @@ -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 diff --git a/build/torch213-cxx11-cu130-x86_64-linux/token_permute/torch_impl.py b/build/torch213-cxx11-cu130-x86_64-linux/token_permute/torch_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..7f06204e57cb6fc96ccb9e10ac9f798ac2d0e7a4 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/token_permute/torch_impl.py @@ -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", +] diff --git a/build/torch213-cxx11-cu130-x86_64-linux/utils/__init__.py b/build/torch213-cxx11-cu130-x86_64-linux/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..99f6c1e8cca79bfdf04640b8a92602b205a407a0 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/utils/__init__.py @@ -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. +# +################################################################################################# diff --git a/build/torch213-cxx11-cu130-x86_64-linux/utils/checks.py b/build/torch213-cxx11-cu130-x86_64-linux/utils/checks.py new file mode 100644 index 0000000000000000000000000000000000000000..f58db50a581e0bc6b43e8de0fbccd9404beaa1dc --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/utils/checks.py @@ -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, + ) diff --git a/build/torch213-cxx11-cu130-x86_64-linux/utils/device.py b/build/torch213-cxx11-cu130-x86_64-linux/utils/device.py new file mode 100644 index 0000000000000000000000000000000000000000..7997ab474197c78f3d41138f17ae311d6e06b1cf --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/utils/device.py @@ -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 diff --git a/build/torch213-cxx11-cu130-x86_64-linux/utils/dtype.py b/build/torch213-cxx11-cu130-x86_64-linux/utils/dtype.py new file mode 100644 index 0000000000000000000000000000000000000000..c65d9e9e0082c6f14af0b8331216cf26fcb3f28d --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/utils/dtype.py @@ -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] diff --git a/build/torch213-cxx11-cu130-x86_64-linux/utils/environment.py b/build/torch213-cxx11-cu130-x86_64-linux/utils/environment.py new file mode 100644 index 0000000000000000000000000000000000000000..0581aacb1183a5dea228a63859837113429e1a62 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/utils/environment.py @@ -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 diff --git a/build/torch213-cxx11-cu130-x86_64-linux/utils/log.py b/build/torch213-cxx11-cu130-x86_64-linux/utils/log.py new file mode 100644 index 0000000000000000000000000000000000000000..138c3ac8763480c0a381a8af0d913efcfa48b1a1 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/utils/log.py @@ -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) diff --git a/build/torch213-cxx11-cu130-x86_64-linux/utils/tensor.py b/build/torch213-cxx11-cu130-x86_64-linux/utils/tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..32d95acf2eb9a8c6fa29c7d3de3b48b5f14be415 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/utils/tensor.py @@ -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 diff --git a/build/torch213-cxx11-cu130-x86_64-linux/utils/testing.py b/build/torch213-cxx11-cu130-x86_64-linux/utils/testing.py new file mode 100644 index 0000000000000000000000000000000000000000..8ed9de5df24754dc9207a0b39954c8a73d7610a5 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/utils/testing.py @@ -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 diff --git a/build/torch213-cxx11-cu130-x86_64-linux/utils/tuples.py b/build/torch213-cxx11-cu130-x86_64-linux/utils/tuples.py new file mode 100644 index 0000000000000000000000000000000000000000..abf8a862fa4b647aab55fb2982e5dea08609d316 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/utils/tuples.py @@ -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 diff --git a/build/torch213-cxx11-cu130-x86_64-linux/utils/varlen.py b/build/torch213-cxx11-cu130-x86_64-linux/utils/varlen.py new file mode 100644 index 0000000000000000000000000000000000000000..8bc541dbd8e80ddf9923a097ac473b31035174ef --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/utils/varlen.py @@ -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, + ) diff --git a/build/torch213-cxx11-cu130-x86_64-linux/version.py b/build/torch213-cxx11-cu130-x86_64-linux/version.py new file mode 100644 index 0000000000000000000000000000000000000000..fe22ed72b87e6b242b01aae1ffb3de069ae7b0a9 --- /dev/null +++ b/build/torch213-cxx11-cu130-x86_64-linux/version.py @@ -0,0 +1,24 @@ +################################################################################################# +# 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. +# +################################################################################################# + +__version__ = "0.21.7" diff --git a/build/torch213-cxx11-cu132-x86_64-linux/__init__.py b/build/torch213-cxx11-cu132-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9fe10d400adaa8333f0da9c0c9654864659fbc3b --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/__init__.py @@ -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", +] diff --git a/build/torch213-cxx11-cu132-x86_64-linux/_environment.py b/build/torch213-cxx11-cu132-x86_64-linux/_environment.py new file mode 100644 index 0000000000000000000000000000000000000000..caae85bbe3c5eef992e4c6758e44b382382ebe41 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/_environment.py @@ -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", +] diff --git a/build/torch213-cxx11-cu132-x86_64-linux/_libnatten/__init__.py b/build/torch213-cxx11-cu132-x86_64-linux/_libnatten/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9566f8d4e7d9356415bb2951613dc5ada194dbd9 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/_libnatten/__init__.py @@ -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", +] diff --git a/build/torch213-cxx11-cu132-x86_64-linux/_libnatten/torch_wrappers.py b/build/torch213-cxx11-cu132-x86_64-linux/_libnatten/torch_wrappers.py new file mode 100644 index 0000000000000000000000000000000000000000..16e5b62b2bc5de0810f58d82d70d7fb2e3e56e4a --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/_libnatten/torch_wrappers.py @@ -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", +] diff --git a/build/torch213-cxx11-cu132-x86_64-linux/_natten_cuda_3641131.abi3.so b/build/torch213-cxx11-cu132-x86_64-linux/_natten_cuda_3641131.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..e58f96dbb9af844ab3832fa49f278bc3a270cf30 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/_natten_cuda_3641131.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1b9224fcbb1a2b89e14a0c4f3a5b896fa768924fc7070cd922195614826af113 +size 165970616 diff --git a/build/torch213-cxx11-cu132-x86_64-linux/_ops.py b/build/torch213-cxx11-cu132-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..ee2d34306175408457b5e3dca950160c1444f947 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/_ops.py @@ -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}" diff --git a/build/torch213-cxx11-cu132-x86_64-linux/_types.py b/build/torch213-cxx11-cu132-x86_64-linux/_types.py new file mode 100644 index 0000000000000000000000000000000000000000..804980900b74586f5861f473853055c7d9630efd --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/_types.py @@ -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 diff --git a/build/torch213-cxx11-cu132-x86_64-linux/attn_merge.py b/build/torch213-cxx11-cu132-x86_64-linux/attn_merge.py new file mode 100644 index 0000000000000000000000000000000000000000..e71cc799d407df0afa701942ad08b03718ba20bf --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/attn_merge.py @@ -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"] diff --git a/build/torch213-cxx11-cu132-x86_64-linux/backends/__init__.py b/build/torch213-cxx11-cu132-x86_64-linux/backends/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..84b8192a37e048880aba2059a6f47e7d26f6a8ee --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/backends/__init__.py @@ -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", +] diff --git a/build/torch213-cxx11-cu132-x86_64-linux/backends/blackwell_fmha.py b/build/torch213-cxx11-cu132-x86_64-linux/backends/blackwell_fmha.py new file mode 100644 index 0000000000000000000000000000000000000000..9b36e79d8b9242b7c5a20833932721556abd61d5 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/backends/blackwell_fmha.py @@ -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 diff --git a/build/torch213-cxx11-cu132-x86_64-linux/backends/blackwell_fna.py b/build/torch213-cxx11-cu132-x86_64-linux/backends/blackwell_fna.py new file mode 100644 index 0000000000000000000000000000000000000000..5bd6f0ed8dd3da85bd60804e503f2aca27b37e4e --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/backends/blackwell_fna.py @@ -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, + ) diff --git a/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/__init__.py b/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9c2b75f7e5007dd1490e8ff5d4c6413f21a5e39f --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/__init__.py @@ -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) diff --git a/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/checks.py b/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/checks.py new file mode 100644 index 0000000000000000000000000000000000000000..4237d8eba4f2e11f7f7abfc796d53b87801f34a6 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/checks.py @@ -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 diff --git a/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass/__init__.py b/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fdd2d4e18d0d959b16f7c9bd0a8b2f71081845be --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass/__init__.py @@ -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}" + ) diff --git a/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass/backward_knobs.py b/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass/backward_knobs.py new file mode 100644 index 0000000000000000000000000000000000000000..dfb1c10a3f1dfa28597f6eb30e5f925ad4785325 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass/backward_knobs.py @@ -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.") diff --git a/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_backward_128x128.py b/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_backward_128x128.py new file mode 100644 index 0000000000000000000000000000000000000000..af6047ccb59f6923be82d6379b92c39b9bcb7f1a --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_backward_128x128.py @@ -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)), + ], +} diff --git a/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_backward_128x64.py b/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_backward_128x64.py new file mode 100644 index 0000000000000000000000000000000000000000..6a915e9686203919079092eccc743dce88b4107a --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_backward_128x64.py @@ -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)), + ], +} diff --git a/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_backward_64x64.py b/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_backward_64x64.py new file mode 100644 index 0000000000000000000000000000000000000000..8511bd844f34a263d233d4e6a1f6f7c775139e2d --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_backward_64x64.py @@ -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)), + ], +} diff --git a/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_forward_32x128.py b/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_forward_32x128.py new file mode 100644 index 0000000000000000000000000000000000000000..a930aba67df1de878d7e95797bb3b6d2d903e020 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_forward_32x128.py @@ -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)), + ], +} diff --git a/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_forward_64x128.py b/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_forward_64x128.py new file mode 100644 index 0000000000000000000000000000000000000000..1bc1fc78e028a755b755ea8108264eccf0bb9659 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_forward_64x128.py @@ -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)), + ], +} diff --git a/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_forward_64x64.py b/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_forward_64x64.py new file mode 100644 index 0000000000000000000000000000000000000000..c321c5b4e1a96f0cb58894f95dd2826d2d21c950 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass/fna_forward_64x64.py @@ -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)), + ], +} diff --git a/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass_blackwell/__init__.py b/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass_blackwell/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4a8c87c6c6d916f58c533fc2a84c66aff8bfe46b --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass_blackwell/__init__.py @@ -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}" + ) diff --git a/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass_hopper/__init__.py b/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass_hopper/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..70b2e3106003ed583050be356b54e32c9b7c61c0 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/cutlass_hopper/__init__.py @@ -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}" + ) diff --git a/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/flex/__init__.py b/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/flex/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c5027701152d3d5976670fd59dc158331c5cadfa --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/backends/configs/flex/__init__.py @@ -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}" + ) diff --git a/build/torch213-cxx11-cu132-x86_64-linux/backends/flex.py b/build/torch213-cxx11-cu132-x86_64-linux/backends/flex.py new file mode 100644 index 0000000000000000000000000000000000000000..0555a7ebdb62e6e4060fbfe3106bed116b44797b --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/backends/flex.py @@ -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, + ) diff --git a/build/torch213-cxx11-cu132-x86_64-linux/backends/fmha.py b/build/torch213-cxx11-cu132-x86_64-linux/backends/fmha.py new file mode 100644 index 0000000000000000000000000000000000000000..0c6851f8bb362a4c7d45db46e61db1ed1dc61d51 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/backends/fmha.py @@ -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 diff --git a/build/torch213-cxx11-cu132-x86_64-linux/backends/fna.py b/build/torch213-cxx11-cu132-x86_64-linux/backends/fna.py new file mode 100644 index 0000000000000000000000000000000000000000..747ecb1f7dcfe032e02e4fa76b8a3a9b70e96bd0 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/backends/fna.py @@ -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, + ) diff --git a/build/torch213-cxx11-cu132-x86_64-linux/backends/hopper_fmha.py b/build/torch213-cxx11-cu132-x86_64-linux/backends/hopper_fmha.py new file mode 100644 index 0000000000000000000000000000000000000000..251b50dc34598111f288ed3ac3d2b414ea23f253 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/backends/hopper_fmha.py @@ -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 diff --git a/build/torch213-cxx11-cu132-x86_64-linux/backends/hopper_fna.py b/build/torch213-cxx11-cu132-x86_64-linux/backends/hopper_fna.py new file mode 100644 index 0000000000000000000000000000000000000000..17bb6d5952675043b66b59fcf92ea9eea54facf9 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/backends/hopper_fna.py @@ -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, + ) diff --git a/build/torch213-cxx11-cu132-x86_64-linux/backends/reference.py b/build/torch213-cxx11-cu132-x86_64-linux/backends/reference.py new file mode 100644 index 0000000000000000000000000000000000000000..d0f34b41a8c9b1a1051360b23d7e20573931bb27 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/backends/reference.py @@ -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, + ) diff --git a/build/torch213-cxx11-cu132-x86_64-linux/context.py b/build/torch213-cxx11-cu132-x86_64-linux/context.py new file mode 100644 index 0000000000000000000000000000000000000000..1793ac984bbe5be383bec2c07b54c770c9969b42 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/context.py @@ -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) diff --git a/build/torch213-cxx11-cu132-x86_64-linux/functional.py b/build/torch213-cxx11-cu132-x86_64-linux/functional.py new file mode 100644 index 0000000000000000000000000000000000000000..944a0c2a7786de250fd5031e933133ea676d2a79 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/functional.py @@ -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, + ) diff --git a/build/torch213-cxx11-cu132-x86_64-linux/metadata.json b/build/torch213-cxx11-cu132-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..a3003e1fad46a43bb8ad0e66bebde3001869b036 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/metadata.json @@ -0,0 +1,84 @@ +{ + "name": "natten", + "id": "_natten_cuda_3641131", + "version": 1, + "license": "MIT", + "upstream": "https://github.com/SHI-Labs/NATTEN", + "python-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "10.0", + "10.0a", + "12.0", + "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": "G5Ik/LsaK4nhSgxPOluJb6dokk/HBwzZIhlWFIJq8RM=", + "_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 + } + } +} \ No newline at end of file diff --git a/build/torch213-cxx11-cu132-x86_64-linux/metadata.json.sigstore b/build/torch213-cxx11-cu132-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..5e12c646d4360fb99ed5d8ccabbb4c5a25a2cf5b --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSjCCBtGgAwIBAgIUG7nkGVVzBnBeb+TXuXdEBR8WJucwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNzI5MDkyMTMzWhcNMjYwNzI5MDkzMTMzWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEW7gmWzajGRVli7U8Lq5hgYkXYg1RQ0WTudAEDQIpmWG96FWvn8HXxhnXu1oWGrwkZp4uFOoCK+BjlTqqRAaPFqOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUcxJllLepJ4d9E3Av6/+pfTlsz3cwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoMjBhNjEzYzEwZDliMWM5Y2I0M2E1MWU4NTBlOWZhMDQ2OWE5NzE5OTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoMjBhNjEzYzEwZDliMWM5Y2I0M2E1MWU4NTBlOWZhMDQ2OWE5NzE5OTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoMjBhNjEzYzEwZDliMWM5Y2I0M2E1MWU4NTBlOWZhMDQ2OWE5NzE5OTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDIwYTYxM2MxMGQ5YjFjOWNiNDNhNTFlODUwZTlmYTA0NjlhOTcxOTkwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzA0MTgyNTI4NjkvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABn60t40QAAAQDAEcwRQIgHwMbAAN+4nG/LCaziEn1PEOlkAJcXMNU6t3J7ugRtYsCIQCVXk/DXnBCMisSZmCT7axL1yS9acoKH6hWcAlpoMcFezAKBggqhkjOPQQDAwNnADBkAjASe7tZy6z3+GD+ZtvB3LwNlMS9X7XN4kKeGz8QHPjLKkU1RGnvupcwfEKmiNkcfJICMHjLgJ+YySnbrP1/t7cOdUoAZ3CoEQfoNDdPlwdv30T7V1MuT4wa1aZF5Zs5lzgbTg=="}, "tlogEntries":[{"logIndex":"2280149214", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1785316893", "inclusionPromise":{"signedEntryTimestamp":"MEUCIFMR4BeiU90wycs/pJgFjnts3RRpPSI06uXnOnuVYnkeAiEAsFQHFsYXYIgSCJhGj4XdWvhWMbd0tA/lBWqTQl9xcCM="}, "inclusionProof":{"logIndex":"2158244952", "rootHash":"EgFZrdyiHsY7+eiEgWLsd08ST3QmRI4XHrDL1PJHNBY=", "treeSize":"2158244954", "hashes":["NW8MBqATPXWci8deQ7UX4MQf3vB2gsXObSsAZrlQxr8=", "6ZrpvK63FOjS3yqMn+ILayFY7mvmnSGS8d2TEhtXrR0=", "NQRnGfc26u6ot2Ztj7yk9Wrroh4XAldKtbHo9z3wnjM=", "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\n2158244954\nEgFZrdyiHsY7+eiEgWLsd08ST3QmRI4XHrDL1PJHNBY=\n\n— rekor.sigstore.dev wNI9ajBFAiEAmbofupOEAS/B5rPajjJ0eRYGy+M7vA7gKKT5jRM8x70CIF+GLNxp6ixFNz3pDJeXlWOPbTRPuFeD1r/JfhJn0k93\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiIzZjVlNjJmOGExOGViMjdiODM0YzQ1YWI0YjUxYjk2ZWUyYjkxNjI4MDkwZGFjN2FmNDliOGU4MjljZjA5ZTFmIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FUUNJRnRENmtpYVdjcmswY0NKYkhXNzBJc1BaTXlGYzg5Z1pJNEkwVVFBa05icUFpQnRxSkF5MWNHU0l1alFBTlhpUTJMdlFNc0gxVlZYRVhUY25PdENzNDNka3c9PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRha05EUW5SSFowRjNTVUpCWjBsVlJ6ZHVhMGRXVm5wQ2JrSmxZaXRVV0hWWVpFVkNVamhYU25WamQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDU2U1RWTlJHdDVUVlJOZWxkb1kwNU5hbGwzVG5wSk5VMUVhM3BOVkUxNlYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZYTjJkdFYzcGhha2RTVm14cE4xVTRUSEUxYUdkWmExaFpaekZTVVRCWFZIVmtRVVVLUkZGSmNHMVhSemsyUmxkMmJqaElXSGhvYmxoMU1XOVhSM0ozYTFwd05IVkdUMjlEU3l0Q2FteFVjWEZTUVdGUVJuRlBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZqZUVwc0NteE1aWEJLTkdRNVJUTkJkall2SzNCbVZHeHplak5qZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU5ha0pvVG1wRmVsbDZSWGRhUkd4cFRWZE5OVmt5U1RCTk1rVXhDazFYVlRST1ZFSnNUMWRhYUUxRVVUSlBWMFUxVG5wRk5VOVVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDFxUW1oT2FrVjZXWHBGZDFwRWJHbE5WMDAxV1RKSk1FMHlSVEZOVjFVMFRsUkNiRTlYV21nS1RVUlJNazlYUlRWT2VrVTFUMVJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDFxUW1oT2FrVjZXWHBGZDFwRWJHbE5WMDAxV1RKSk1FMHlSVEVLVFZkVk5FNVVRbXhQVjFwb1RVUlJNazlYUlRWT2VrVTFUMVJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUkpkMWxVV1hnS1RUSk5lRTFIVVRWWmFrWnFUMWRPYVU1RVRtaE9WRVpzVDBSVmQxcFViRzFaVkVFd1RtcHNhRTlVWTNoUFZHdDNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWtFd1RWUm5lVTVVU1RST2FtdDJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNDJNSFEwTUZGQlFVRlJSQXBCUldOM1VsRkpaMGgzVFdKQlFVNHJORzVITDB4RFlYcHBSVzR4VUVWUGJHdEJTbU5ZVFU1Vk5uUXpTamQxWjFKMFdYTkRTVkZEVmxockwwUllia0pEQ2sxcGMxTmFiVU5VTjJGNFRERjVVemxoWTI5TFNEWm9WMk5CYkhCdlRXTkdaWHBCUzBKblozRm9hMnBQVUZGUlJFRjNUbTVCUkVKclFXcEJVMlUzZEZvS2VUWjZNeXRIUkN0YWRIWkNNMHgzVG14TlV6bFlOMWhPTkd0TFpVZDZPRkZJVUdwTVMydFZNVkpIYm5aMWNHTjNaa1ZMYldsT2EyTm1Ta2xEVFVocVRBcG5TaXRaZVZOdVluSlFNUzkwTjJOUFpGVnZRVm96UTI5RlVXWnZUa1JrVUd4M1pIWXpNRlEzVmpGTmRWUTBkMkV4WVZwR05WcHpOV3g2WjJKVVp6MDlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgqNSIt5p3WfQfKu6ZkHqJOTKcHhiKMKeK7LOSF6VtgDUCFQCEy/HemoCpeH7+ttDLSuIZqiRvCRgPMjAyNjA3MjkwOTIxMzNaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHbMIIB1wIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDcyOTA5MjEzM1owLwYJKoZIhvcNAQkEMSIEIPaor8WRyLOfCsJ8+iIwBllrDY8MG7yJ+l81NoBHj1keMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRnMGUCMEpydCdREDBg1YgRDRPHBBPDw1SF22JsCWTevdNEP4svdGGimJe6DtWwlntyVCCHVQIxAKsDcPwGODoWuzvS3/AktHWe1t8Yzi5OKLaB4EU2ho4CN9qq4QQVmqkJREjODuvG6A=="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"P15i+KGOsnuDTEWrS1G5buK5FigJDax69JuOgpzwnh8="}, "signature":"MEQCIFtD6kiaWcrk0cCJbHW70IsPZMyFc89gZI4I0UQAkNbqAiBtqJAy1cGSIujQANXiQ2LvQMsH1VVXEXTcnOtCs43dkw=="}} \ No newline at end of file diff --git a/build/torch213-cxx11-cu132-x86_64-linux/modules.py b/build/torch213-cxx11-cu132-x86_64-linux/modules.py new file mode 100644 index 0000000000000000000000000000000000000000..222e384f2151c6a32275abcdb5000d5b8f1379d4 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/modules.py @@ -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, + ) diff --git a/build/torch213-cxx11-cu132-x86_64-linux/natten/__init__.py b/build/torch213-cxx11-cu132-x86_64-linux/natten/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/natten/__init__.py @@ -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"))) diff --git a/build/torch213-cxx11-cu132-x86_64-linux/token_permute/__init__.py b/build/torch213-cxx11-cu132-x86_64-linux/token_permute/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..381de37e4e7d4d2e6158668fbcc2ca52f7b388a3 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/token_permute/__init__.py @@ -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", +] diff --git a/build/torch213-cxx11-cu132-x86_64-linux/token_permute/cutlass_impl.py b/build/torch213-cxx11-cu132-x86_64-linux/token_permute/cutlass_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..49ce678f9052679448f843922f16bdec5d8796be --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/token_permute/cutlass_impl.py @@ -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 diff --git a/build/torch213-cxx11-cu132-x86_64-linux/token_permute/frontend.py b/build/torch213-cxx11-cu132-x86_64-linux/token_permute/frontend.py new file mode 100644 index 0000000000000000000000000000000000000000..720aded02e93e31f19a91efbda82cac3db15e4d2 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/token_permute/frontend.py @@ -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 diff --git a/build/torch213-cxx11-cu132-x86_64-linux/token_permute/torch_impl.py b/build/torch213-cxx11-cu132-x86_64-linux/token_permute/torch_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..7f06204e57cb6fc96ccb9e10ac9f798ac2d0e7a4 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/token_permute/torch_impl.py @@ -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", +] diff --git a/build/torch213-cxx11-cu132-x86_64-linux/utils/__init__.py b/build/torch213-cxx11-cu132-x86_64-linux/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..99f6c1e8cca79bfdf04640b8a92602b205a407a0 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/utils/__init__.py @@ -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. +# +################################################################################################# diff --git a/build/torch213-cxx11-cu132-x86_64-linux/utils/checks.py b/build/torch213-cxx11-cu132-x86_64-linux/utils/checks.py new file mode 100644 index 0000000000000000000000000000000000000000..f58db50a581e0bc6b43e8de0fbccd9404beaa1dc --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/utils/checks.py @@ -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, + ) diff --git a/build/torch213-cxx11-cu132-x86_64-linux/utils/device.py b/build/torch213-cxx11-cu132-x86_64-linux/utils/device.py new file mode 100644 index 0000000000000000000000000000000000000000..7997ab474197c78f3d41138f17ae311d6e06b1cf --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/utils/device.py @@ -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 diff --git a/build/torch213-cxx11-cu132-x86_64-linux/utils/dtype.py b/build/torch213-cxx11-cu132-x86_64-linux/utils/dtype.py new file mode 100644 index 0000000000000000000000000000000000000000..c65d9e9e0082c6f14af0b8331216cf26fcb3f28d --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/utils/dtype.py @@ -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] diff --git a/build/torch213-cxx11-cu132-x86_64-linux/utils/environment.py b/build/torch213-cxx11-cu132-x86_64-linux/utils/environment.py new file mode 100644 index 0000000000000000000000000000000000000000..0581aacb1183a5dea228a63859837113429e1a62 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/utils/environment.py @@ -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 diff --git a/build/torch213-cxx11-cu132-x86_64-linux/utils/log.py b/build/torch213-cxx11-cu132-x86_64-linux/utils/log.py new file mode 100644 index 0000000000000000000000000000000000000000..138c3ac8763480c0a381a8af0d913efcfa48b1a1 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/utils/log.py @@ -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) diff --git a/build/torch213-cxx11-cu132-x86_64-linux/utils/tensor.py b/build/torch213-cxx11-cu132-x86_64-linux/utils/tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..32d95acf2eb9a8c6fa29c7d3de3b48b5f14be415 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/utils/tensor.py @@ -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 diff --git a/build/torch213-cxx11-cu132-x86_64-linux/utils/testing.py b/build/torch213-cxx11-cu132-x86_64-linux/utils/testing.py new file mode 100644 index 0000000000000000000000000000000000000000..8ed9de5df24754dc9207a0b39954c8a73d7610a5 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/utils/testing.py @@ -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 diff --git a/build/torch213-cxx11-cu132-x86_64-linux/utils/tuples.py b/build/torch213-cxx11-cu132-x86_64-linux/utils/tuples.py new file mode 100644 index 0000000000000000000000000000000000000000..abf8a862fa4b647aab55fb2982e5dea08609d316 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/utils/tuples.py @@ -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 diff --git a/build/torch213-cxx11-cu132-x86_64-linux/utils/varlen.py b/build/torch213-cxx11-cu132-x86_64-linux/utils/varlen.py new file mode 100644 index 0000000000000000000000000000000000000000..8bc541dbd8e80ddf9923a097ac473b31035174ef --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/utils/varlen.py @@ -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, + ) diff --git a/build/torch213-cxx11-cu132-x86_64-linux/version.py b/build/torch213-cxx11-cu132-x86_64-linux/version.py new file mode 100644 index 0000000000000000000000000000000000000000..fe22ed72b87e6b242b01aae1ffb3de069ae7b0a9 --- /dev/null +++ b/build/torch213-cxx11-cu132-x86_64-linux/version.py @@ -0,0 +1,24 @@ +################################################################################################# +# 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. +# +################################################################################################# + +__version__ = "0.21.7"