xformers 0.0.35.dev1146__py3-none-any.whl

xformers/__init__.py

@@ -0,0 +1,73 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+#
+# This source code is licensed under the BSD license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+import os
+
+import torch
+
+from . import _cpp_lib
+from .checkpoint import (  # noqa: E402, F401
+    checkpoint,
+    get_optimal_checkpoint_policy,
+    list_operators,
+    selective_checkpoint_wrapper,
+)
+
+try:
+    from .version import __version__  # noqa: F401
+except ImportError:
+    __version__ = "0.0.0"
+
+
+logger = logging.getLogger("xformers")
+
+_has_cpp_library: bool = _cpp_lib._has_cpp_library
+
+_is_opensource: bool = True
+
+
+def compute_once(func):
+    value = None
+
+    def func_wrapper():
+        nonlocal value
+        if value is None:
+            value = func()
+        return value
+
+    return func_wrapper
+
+
+@compute_once
+def _is_triton_available():
+    if os.environ.get("XFORMERS_ENABLE_TRITON", "0") == "1":
+        return True
+    if not torch.cuda.is_available():
+        return False
+    if os.environ.get("XFORMERS_FORCE_DISABLE_TRITON", "0") == "1":
+        return False
+    # We have many errors on V100 with recent triton versions
+    # Let's just drop support for triton kernels below A100
+    if torch.cuda.get_device_capability("cuda") < (8, 0):
+        return False
+    try:
+        import triton  # noqa
+
+        return True
+    except (ImportError, AttributeError):
+        logger.warning(
+            "A matching Triton is not available, some optimizations will not be enabled",
+            exc_info=True,
+        )
+        return False
+
+
+@compute_once
+def get_python_lib():
+    return torch.library.Library("xformers_python", "DEF")
+
+
+# end of file

xformers/_cpp_lib.py

@@ -0,0 +1,138 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+#
+# This source code is licensed under the BSD license found in the
+# LICENSE file in the root directory of this source tree.
+
+import dataclasses
+import json
+import logging
+import os
+import platform
+from typing import Any, Dict, Optional
+
+import torch
+
+logger = logging.getLogger("xformers")
+
+UNAVAILABLE_FEATURES_MSG = "  SwiGLU, sparse and more won't be available."
+
+
+@dataclasses.dataclass
+class _BuildInfo:
+    metadata: Dict[str, Any]
+
+    @property
+    def cuda_version(self) -> Optional[int]:
+        return self.metadata["version"]["cuda"]
+
+    @property
+    def hip_version(self) -> Optional[int]:
+        return self.metadata["version"]["hip"]
+
+    @property
+    def torch_version(self) -> str:
+        return self.metadata["version"]["torch"]
+
+    @property
+    def python_version(self) -> str:
+        return self.metadata["version"]["python"]
+
+    @property
+    def flash_version(self) -> str:
+        return self.metadata["version"].get("flash", "0.0.0")
+
+    @property
+    def use_torch_flash(self) -> bool:
+        return self.metadata["version"].get("use_torch_flash", False)
+
+    @property
+    def build_env(self) -> Dict[str, Any]:
+        return self.metadata["env"]
+
+
+class xFormersWasNotBuiltException(Exception):
+    def __str__(self) -> str:
+        return (
+            "Need to compile C++ extensions to use all xFormers features.\n"
+            "    Please install xformers properly "
+            "(see https://github.com/facebookresearch/xformers#installing-xformers)\n"
+            + UNAVAILABLE_FEATURES_MSG
+        )
+
+
+class xFormersInvalidLibException(Exception):
+    def __init__(self, build_info: Optional[_BuildInfo]) -> None:
+        self.build_info = build_info
+
+    def __str__(self) -> str:
+        if self.build_info is None:
+            msg = "xFormers was built for a different version of PyTorch or Python."
+        else:
+            msg = f"""xFormers was built for:
+    PyTorch {self.build_info.torch_version} with CUDA {self.build_info.cuda_version} (you have {torch.__version__})
+    Python  {self.build_info.python_version} (you have {platform.python_version()})"""
+        return (
+            "xFormers can't load C++/CUDA extensions. "
+            + msg
+            + "\n  Please reinstall xformers "
+            "(see https://github.com/facebookresearch/xformers#installing-xformers)\n"
+            + UNAVAILABLE_FEATURES_MSG
+        )
+
+
+def _register_extensions():
+    import importlib
+    import os
+
+    import torch
+
+    # load the custom_op_library and register the custom ops
+    lib_dir = os.path.dirname(__file__)
+
+    loader_details = (
+        importlib.machinery.ExtensionFileLoader,
+        importlib.machinery.EXTENSION_SUFFIXES,
+    )
+
+    extfinder = importlib.machinery.FileFinder(lib_dir, loader_details)
+    if torch.version.hip and not hasattr(torch.version, "git_version"):
+        ext_specs = extfinder.find_spec("_C_hip")
+    else:
+        ext_specs = extfinder.find_spec("_C")
+    cpp_lib_json = os.path.join(lib_dir, "cpp_lib.json")
+    if ext_specs is None and not os.path.exists(cpp_lib_json):
+        # Nothing was compiled, and nothing was meant to be: this is the
+        # pure-Python build. Everything that needs a kernel now lives in
+        # separate packages, such as mslk.
+        return None
+    if ext_specs is None:
+        raise xFormersWasNotBuiltException()
+    with open(cpp_lib_json, "r") as fp:
+        build_metadata = _BuildInfo(json.load(fp))
+    try:
+        torch.ops.load_library(ext_specs.origin)
+    except OSError as exc:
+        raise xFormersInvalidLibException(build_metadata) from exc
+    return build_metadata
+
+
+_cpp_library_load_exception = None
+_build_metadata: Optional[_BuildInfo] = None
+_has_cpp_library: bool = False
+
+try:
+    _build_metadata = _register_extensions()
+    _has_cpp_library = _build_metadata is not None
+except (xFormersInvalidLibException, xFormersWasNotBuiltException) as e:
+    ENV_VAR_FOR_DETAILS = "XFORMERS_MORE_DETAILS"
+    if os.environ.get(ENV_VAR_FOR_DETAILS, False):
+        logger.warning(f"WARNING[XFORMERS]: {e}", exc_info=e)
+    else:
+        logger.warning(
+            f"WARNING[XFORMERS]: {e}\n  Set {ENV_VAR_FOR_DETAILS}=1 for more details"
+        )
+    _cpp_library_load_exception = e
+
+_built_with_cuda = (
+    _build_metadata is not None and _build_metadata.cuda_version is not None
+)

xformers/_deprecation_warning.py

@@ -0,0 +1,12 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+#
+# This source code is licensed under the BSD license found in the
+# LICENSE file in the root directory of this source tree.
+
+import warnings
+
+
+def deprecated_function(self):
+    name = repr(self)  # self.__name__
+    msg = f"{name} is deprecated and is not maintained anymore. It might be removed in a future version of xFormers"
+    warnings.warn(msg, FutureWarning, stacklevel=2)

xformers/attn_bias_utils.py

@@ -0,0 +1,12 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+#
+# This source code is licensed under the BSD license found in the
+# LICENSE file in the root directory of this source tree.
+# flake8: noqa
+
+from mslk.attention.fmha.attn_bias_utils import (
+    create_attn_bias,
+    pack_kv_cache,
+    ref_attention,
+    ref_attention_bmhk,
+)

xformers/benchmarks/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+#
+# This source code is licensed under the BSD license found in the
+# LICENSE file in the root directory of this source tree.

xformers/benchmarks/benchmark_attn_decoding.py

@@ -0,0 +1,467 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+#
+# This source code is licensed under the BSD license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Run with --omit-baselines to skip slow baselines.
+# See other CLI arguments in benchmark_main_helper in utils.py.
+
+import sys
+from typing import Any, Dict, Type
+
+import pytest
+import torch
+
+import xformers.ops as xops
+from xformers.attn_bias_utils import create_attn_bias
+from xformers.benchmarks.utils import benchmark_main_helper2, NotSupportedInputError
+
+min_run_time = 0.5
+device = torch.device("cuda")
+
+
+CASES = [
+    dict(
+        B=max(1, 2 ** (16 - i)),
+        Mq=1,
+        Mkv=2**i,
+        Hq=16,
+        Hkv=hkv,
+        K=128,
+        attn_bias_type=xops.fmha.attn_bias.BlockDiagonalCausalWithOffsetPaddedKeysMask,
+    )
+    for i in range(8, 18)
+    for hkv in (1, 2)
+]
+
+
+def quantize_kv_int4(k: torch.Tensor, num_groups: int = 1) -> torch.Tensor:
+    """
+    Auxiliary int4 row quantization function used for benchmarking and tests.
+    Matches the behaviour of torch.ops.llama_cpp.dequantize_int4_cache -
+    quantization parameters (scale and offset) of each row along the last
+    dimension of the tensor are assumed to be packed into two float16 values
+    at the beginning of the row.
+    """
+    # Scale and shift are such that quantization linearly maps int4 values range [0..15]
+    # to input values range min(k)..max(k) individually for every row
+    k = k.reshape(*k.shape[:-1], num_groups, k.shape[-1] // num_groups)
+    # print(f"k_reshape = {k.shape}")
+    max_vals = torch.max(k, dim=-1, keepdim=True).values
+    min_vals = torch.min(k, dim=-1, keepdim=True).values
+    scale_k: torch.Tensor = (max_vals - min_vals) / 15
+    # print(f"scale_k_shape = {scale_k.shape}")
+
+    shift_k = torch.min(k, dim=-1, keepdim=True).values
+    scale_k = scale_k.to(torch.float16)
+    shift_k = shift_k.to(torch.float16)
+    in_bytes = ((k - shift_k.expand(k.shape)) / scale_k.expand(k.shape)) + 0.5
+    in_bytes = in_bytes.to(torch.uint8)
+    in_int4 = in_bytes & 0xF
+    in_int4_packed = in_int4[..., ::2] + (in_int4[..., 1::2] << 4)
+    scale_shift = torch.concat(
+        [scale_k.view(torch.uint8), shift_k.view(torch.uint8)], dim=-1
+    )
+    k_quant = torch.concat(
+        [
+            scale_shift.flatten(start_dim=-2),
+            in_int4_packed.flatten(start_dim=-2),
+        ],
+        dim=-1,
+    ).view(torch.int16)
+    return k_quant
+
+
+class AttentionDecodingBase:
+    OP: Any = None
+
+    def __init__(
+        self,
+        B: int,
+        Mq: int,
+        Mkv: int,
+        Hq: int,
+        Hkv: int,
+        K: int,
+        bw: bool,
+        attn_bias_type,
+    ) -> None:
+        dtype = torch.float16
+        torch.manual_seed(10)
+        self.sub_label = (
+            f"B={B} Mq={Mq} Mkv={Mkv} Hq={Hq} Hkv={Hkv} K={K} TotalBytes="
+            f"{((B * Mkv * Hkv * K * 2) + (B * Mq * Hq * K) + (B * Mq * Hq * K)) * 2}"
+        )
+        self.label = "attn_decoding"
+        self.shapes = (B, Mq, Mkv, Hq, Hkv, K)
+
+        assert Hkv <= Hq
+        assert Hq % Hkv == 0
+        self.q = torch.randn(
+            [B, Mq, Hkv, Hq // Hkv, K], device="cuda", dtype=dtype, requires_grad=bw
+        )
+        self.k = torch.randn(
+            [B, Mkv, Hkv, 1, K], device="cuda", dtype=dtype, requires_grad=bw
+        ).expand(-1, -1, -1, Hq // Hkv, -1)
+        self.v = torch.randn(
+            [B, Mkv, Hkv, 1, K], device="cuda", dtype=dtype, requires_grad=bw
+        ).expand(-1, -1, -1, Hq // Hkv, -1)
+
+        if Hq == Hkv:
+            self.q = self.q[:, :, :, 0]
+            self.k = self.k[:, :, :, 0]
+            self.v = self.v[:, :, :, 0]
+        if Hkv == 1:
+            self.q = self.q[:, :, 0]
+            self.k = self.k[:, :, 0]
+            self.v = self.v[:, :, 0]
+
+        self.attn_bias = create_attn_bias(
+            attn_bias_type,
+            batch_size=B,
+            num_heads=Hq,
+            num_heads_groups=Hq // Hkv,
+            q_len=Mq,
+            kv_len=Mkv,
+            dtype=dtype,
+            device=device,
+            requires_grad=False,
+            fmt="BMHK",
+            op=self.OP,
+        )
+
+        if isinstance(
+            self.attn_bias,
+            xops.fmha.attn_bias.BlockDiagonalCausalWithOffsetPaddedKeysMask,
+        ):
+            self.q = self.q.view(1, -1, *self.q.shape[2:])
+            self.k = self.k.view(1, -1, *self.k.shape[2:])
+            self.v = self.v.view(1, -1, *self.v.shape[2:])
+
+        if hasattr(self.OP, "not_supported_reasons"):
+            inp = xops.fmha.Inputs(
+                query=self.q, key=self.k, value=self.v, attn_bias=self.attn_bias
+            )
+            not_supported_reasons = self.OP.not_supported_reasons(inp)
+            if not_supported_reasons:
+                raise NotSupportedInputError(not_supported_reasons)
+
+    def get_inputs(self):
+        inp = xops.fmha.Inputs(
+            query=self.q, key=self.k, value=self.v, attn_bias=self.attn_bias
+        )
+        return inp
+
+    def fw(self) -> None:
+        try:
+            xops.memory_efficient_attention_forward(
+                self.q, self.k, self.v, op=self.OP, attn_bias=self.attn_bias
+            )
+        except (RuntimeError, ValueError) as e:
+            print(f"Runtime error: {e}")
+
+
+class AttentionDecodingCUTLASS(AttentionDecodingBase):
+    OP = xops.fmha.cutlass.FwOp
+
+
+class AttentionDecodingCK(AttentionDecodingBase):
+    OP = xops.fmha.ck.FwOp
+
+    def __init__(
+        self,
+        B: int,
+        Mq: int,
+        Mkv: int,
+        Hq: int,
+        Hkv: int,
+        K: int,
+        bw: bool,
+        attn_bias_type,
+    ) -> None:
+        dtype = torch.float16
+        torch.manual_seed(10)
+        self.sub_label = (
+            f"B={B} Mq={Mq} Mkv={Mkv} Hq={Hq} Hkv={Hkv} K={K} TotalBytes="
+            f"{((B * Mkv * Hkv * K * 2) + (B * Mq * Hq * K) + (B * Mq * Hq * K)) * 2}"
+        )
+        self.label = "attn_decoding"
+        self.shapes = (B, Mq, Mkv, Hq, Hkv, K)
+
+        assert Hkv <= Hq
+        assert Hq % Hkv == 0
+        self.q = torch.randn(
+            [B, Mq, Hkv, Hq // Hkv, K], device="cuda", dtype=dtype, requires_grad=bw
+        )
+        self.k = torch.randn(
+            [B, Mkv, Hkv, 1, K], device="cuda", dtype=dtype, requires_grad=bw
+        ).expand(-1, -1, -1, Hq // Hkv, -1)
+        self.v = torch.randn(
+            [B, Mkv, Hkv, 1, K], device="cuda", dtype=dtype, requires_grad=bw
+        ).expand(-1, -1, -1, Hq // Hkv, -1)
+
+        if Hq == Hkv:
+            self.q = self.q[:, :, :, 0]
+            self.k = self.k[:, :, :, 0]
+            self.v = self.v[:, :, :, 0]
+
+        self.attn_bias = create_attn_bias(
+            attn_bias_type,
+            batch_size=B,
+            num_heads=Hq,
+            num_heads_groups=Hq // Hkv,
+            q_len=Mq,
+            kv_len=Mkv,
+            dtype=dtype,
+            device=device,
+            requires_grad=False,
+            fmt="BMHK",
+            op=self.OP,
+        )
+
+        if isinstance(
+            self.attn_bias,
+            xops.fmha.attn_bias.BlockDiagonalCausalWithOffsetPaddedKeysMask,
+        ):
+            self.q = self.q.view(1, -1, *self.q.shape[2:])
+            self.k = self.k.view(1, -1, *self.k.shape[2:])
+            self.v = self.v.view(1, -1, *self.v.shape[2:])
+
+        if hasattr(self.OP, "not_supported_reasons"):
+            inp = xops.fmha.Inputs(
+                query=self.q, key=self.k, value=self.v, attn_bias=self.attn_bias
+            )
+            not_supported_reasons = self.OP.not_supported_reasons(inp)
+            if not_supported_reasons:
+                raise NotSupportedInputError(not_supported_reasons)
+
+
+class AttentionDecodingSplitKV(AttentionDecodingBase):
+    OP = xops.fmha.triton_splitk.FwOp
+
+
+class AttentionDecodingCKSplitKV(AttentionDecodingBase):
+    OP = xops.fmha.ck_splitk.FwOp
+
+
+class AttentionDecodingSplitInt4KV(AttentionDecodingBase):
+    OP = xops.fmha.triton_splitk.FwOp
+
+    def __init__(
+        self,
+        B: int,
+        Mq: int,
+        Mkv: int,
+        Hq: int,
+        Hkv: int,
+        K: int,
+        bw: bool,
+        attn_bias_type,
+    ) -> None:
+        # super(AttentionDecodingSplitInt4KV, self).__init__(B, Mq, Mkv, Hq, Hkv, K, bw, attn_bias_type)
+        dtype = torch.float16
+        torch.manual_seed(10)
+        self.sub_label = (
+            f"B={B} Mq={Mq} Mkv={Mkv} Hq={Hq} Hkv={Hkv} K={K} TotalBytes="
+            f"{((B * Mkv * Hkv * K * 2) + (B * Mq * Hq * K) + (B * Mq * Hq * K)) * 2}"
+        )
+        self.label = "attn_decoding"
+        self.shapes = (B, Mq, Mkv, Hq, Hkv, K)
+
+        assert Hkv <= Hq
+        assert Hq % Hkv == 0
+        self.q = torch.randn(
+            [B, Mq, Hkv, Hq // Hkv, K], device="cuda", dtype=dtype, requires_grad=bw
+        )
+        self.k = torch.randn(
+            [B, Mkv, Hkv, 1, K], device="cuda", dtype=dtype, requires_grad=bw
+        )
+        self.v = torch.randn(
+            [B, Mkv, Hkv, 1, K], device="cuda", dtype=dtype, requires_grad=bw
+        )
+
+        num_groups = 1
+        self.k = (
+            quantize_kv_int4(self.k, num_groups=num_groups)
+            .contiguous()
+            .view(torch.int32)
+        ).expand(-1, -1, -1, Hq // Hkv, -1)
+        self.v = (
+            quantize_kv_int4(self.v, num_groups=num_groups)
+            .contiguous()
+            .view(torch.int32)
+        ).expand(-1, -1, -1, Hq // Hkv, -1)
+
+        if Hq == Hkv:
+            self.q = self.q[:, :, :, 0]
+            self.k = self.k[:, :, :, 0]
+            self.v = self.v[:, :, :, 0]
+        if Hkv == 1:
+            self.q = self.q[:, :, 0]
+            self.k = self.k[:, :, 0]
+            self.v = self.v[:, :, 0]
+
+        self.attn_bias = create_attn_bias(
+            attn_bias_type,
+            batch_size=B,
+            num_heads=Hq,
+            num_heads_groups=Hq // Hkv,
+            q_len=Mq,
+            kv_len=Mkv,
+            dtype=dtype,
+            device=device,
+            requires_grad=False,
+            fmt="BMHK",
+            op=self.OP,
+        )
+
+        if isinstance(
+            self.attn_bias,
+            xops.fmha.attn_bias.BlockDiagonalCausalWithOffsetPaddedKeysMask,
+        ):
+            self.q = self.q.view(1, -1, *self.q.shape[2:])
+            self.k = self.k.view(1, -1, *self.k.shape[2:])
+            self.v = self.v.view(1, -1, *self.v.shape[2:])
+
+        if hasattr(self.OP, "not_supported_reasons"):
+            inp = xops.fmha.Inputs(
+                query=self.q, key=self.k, value=self.v, attn_bias=self.attn_bias
+            )
+            not_supported_reasons = self.OP.not_supported_reasons(inp)
+            if not_supported_reasons:
+                raise NotSupportedInputError(not_supported_reasons)
+
+
+class AttentionDecodingPyTorchRepeat(AttentionDecodingBase):
+    def fw(self) -> None:
+        B, Mq, Mkv, Hq, Hkv, K = self.shapes
+        scale = 1 / K**0.5
+        q = self.q.reshape([B, Mq, -1, K]).permute(0, 2, 1, 3)
+        k = self.k.reshape([B, Mkv, -1, K]).permute(0, 2, 1, 3)
+        v = self.v.reshape([B, Mkv, -1, K]).permute(0, 2, 1, 3)
+        attn = (q @ k.transpose(-1, -2) * scale).softmax(-1)
+        return attn @ v
+
+
+BENCHMARKS: Dict[str, Type[AttentionDecodingBase]] = {
+    "pytorch": AttentionDecodingPyTorchRepeat,
+}
+
+if torch.version.cuda:
+    BENCHMARKS["cutlass"] = AttentionDecodingCUTLASS
+
+if torch.version.hip:
+    BENCHMARKS.update(
+        {
+            "ck": AttentionDecodingCK,
+            "ck_splitK": AttentionDecodingCKSplitKV,
+        }
+    )
+
+
+if (sys.version_info.major, sys.version_info.minor) >= (3, 9):
+    BENCHMARKS["triton_splitK"] = AttentionDecodingSplitKV
+    BENCHMARKS["triton_int4KV"] = AttentionDecodingSplitInt4KV
+
+try:
+    import flash_attn
+
+    class AttentionDecodingFlashAttention(AttentionDecodingBase):
+        def fw(self) -> None:
+            q, k, v = self.q, self.k, self.v
+            if q.ndim == 5:
+                B, Mq, H1, H2, K = q.shape
+                B, Mkv, H1, H2, K = k.shape
+                q = q.reshape([B, Mq, H1 * H2, K])
+                k = k[:, :, :, 0]
+                v = v[:, :, :, 0]
+            return flash_attn.flash_attn_func(q, k, v)
+
+    BENCHMARKS[f"flash-attention@{flash_attn.__version__}"] = (
+        AttentionDecodingFlashAttention
+    )
+except ImportError:
+    pass
+
+
+TEST_CASES = [
+    dict(
+        B=max(1, 2 ** (16 - i)),
+        Mq=1,
+        Mkv=2**i,
+        Hq=16,
+        Hkv=hkv,
+        K=128,
+        attn_bias_type=None,
+    )
+    for i in range(8, 18)
+    for hkv in range(1, 3)
+] + [
+    dict(B=i, Mq=1, Mkv=4097, Hq=8, Hkv=1, K=128, attn_bias_type=None)
+    for i in [2, 4, 8, 16, 32, 64, 128]
+]
+
+
+def get_benchmark_names():
+    decoder_names = list(BENCHMARKS.keys())
+    decoder_names.remove("pytorch")
+    return decoder_names
+
+
+# tests to verify correctness of each decoder implementation
+@pytest.mark.parametrize(
+    "name, case",
+    [(name, case) for name in get_benchmark_names() for case in TEST_CASES],
+)
+def test_flash_attention_decoder(name, case):
+    baseline = AttentionDecodingPyTorchRepeat(
+        case["B"],
+        case["Mq"],
+        case["Mkv"],
+        case["Hq"],
+        case["Hkv"],
+        case["K"],
+        False,
+        case["attn_bias_type"],
+    )
+    if name == "ck-decoder" and case["Mkv"] >= 2**14:
+        pytest.skip("ck-decoder does not support Mkv >= 16K")
+
+    baseline_out = baseline.fw()
+    inputs = baseline.get_inputs()
+    decoder = BENCHMARKS[name]
+
+    assert name in ["ck_splitK", "ck", "triton_splitK", "triton_int4KV"]
+    decoder_output, ctx = decoder.OP.apply(inputs, False)
+
+    q, k, v = inputs.get_qkv_in_bmghk()
+    B, M, G, H, Kq = q.shape
+    mqa_swap_seqlen_head = False
+    if k.shape[3] > 1 and k.stride(3) == 0 and v.stride(3) == 0:
+        mqa_swap_seqlen_head = True
+    if mqa_swap_seqlen_head:
+        decoder_output = (
+            decoder_output.reshape(B, -1, M, Kq).transpose(1, 2).contiguous()
+        )
+    else:
+        decoder_output = decoder_output.reshape(B, H * G, -1, Kq).contiguous()
+
+    decoder_output = decoder_output.transpose(2, 1).contiguous()
+    torch.testing.assert_close(decoder_output, baseline_out, atol=1e-2, rtol=0)
+
+
+def main() -> None:
+    """
+    run performance benchmark
+    """
+    benchmark_main_helper2(
+        "attn_decoding",
+        fw=True,
+        cases=CASES,
+        functions=BENCHMARKS,
+        min_run_time=min_run_time,
+    )
+
+
+if __name__ == "__main__":
+    main()  # pragma: no cover