xformers 0.0.23.dev644__tar.gz → 0.0.23.dev648__tar.gz

{xformers-0.0.23.dev644/xformers.egg-info → xformers-0.0.23.dev648}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xformers
-Version: 0.0.23.dev644
+Version: 0.0.23.dev648
 Summary: XFormers: A collection of composable Transformer building blocks.
 Home-page: https://facebookresearch.github.io/xformers/
 Author: Facebook AI Research

{xformers-0.0.23.dev644 → xformers-0.0.23.dev648}/tests/test_indexing.py

@@ -9,13 +9,15 @@ import pytest
 import torch
 
 import xformers.ops as xops
+from xformers.ops import indexing
 
 from .utils import assert_allclose
 
-cuda_only = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
 
-
-@cuda_only
+@pytest.mark.skipif(
+    not indexing.ScaledIndexAddFw.is_available(), reason="not available"
+)
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
 @pytest.mark.parametrize("with_scaling", [False, True])
 @pytest.mark.parametrize(
     "out_shape", [(48, 1, 257 * 1536), (48, 257, 1536), (192, 50, 1536)]
@@ -71,7 +73,8 @@ def test_scaled_index_add(out_shape, with_scaling: bool) -> None:
         assert_allclose(v.grad, ref_grads[k], f"{k}.grad", atol=atol, rtol=rtol)  # type: ignore
 
 
-@cuda_only
+@pytest.mark.skipif(not indexing.IndexSelect.is_available(), reason="not available")
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
 @pytest.mark.parametrize("D", [1536])
 @pytest.mark.parametrize("batches", [((48, 25), (192, 50))])
 def test_index_select_cat(D, batches) -> None:

{xformers-0.0.23.dev644 → xformers-0.0.23.dev648}/tests/test_mem_eff_attention.py

@@ -2055,6 +2055,47 @@ def test_forward_gqa(opFW_biasT):
     )
 
 
+@cuda_only
+@pytest.mark.parametrize(
+    "opBW",
+    [
+        fmha.flash.BwOp,
+        fmha.cutlass.BwOp,
+    ],
+)
+def test_backward_gqa(opBW):
+    H = 8
+    B_Mq_Mkv_H_K_Kv = (3, 512, 512, H, 128, 128)
+    dtype = torch.float16
+    query, key, value, attn_bias = create_tensors(
+        *(opBW, "cuda", dtype, type(None), *B_Mq_Mkv_H_K_Kv),
+        attn_bias_requires_grad=False,
+        fmt="BMHK",
+    )
+    op = (fmha.cutlass.FwOp, opBW)
+    key = key[:, :, :1].expand(-1, -1, H, -1)
+    value = value[:, :, :1].expand(-1, -1, H, -1)
+    key.requires_grad_(True)
+    out = fmha.memory_efficient_attention(query, key, value, attn_bias=attn_bias)
+    out_ref = ref_attention_bmhk(query, key, value, attn_bias=attn_bias)
+    assert_allclose(
+        out.float(),
+        out_ref.float(),
+        atol=op[0].ERROR_ATOL[dtype],
+        rtol=op[0].ERROR_RTOL[dtype],
+    )
+    out.backward(query)
+    dk = key.grad
+    key.grad = None
+    out_ref.backward(query)
+    assert_allclose(
+        dk.float(),
+        key.grad.float(),
+        atol=op[1].ERROR_ATOL[dtype],
+        rtol=op[1].ERROR_RTOL[dtype],
+    )
+
+
 @cuda_only
 @pytest.mark.parametrize("opFW", [op for op in ALL_FW_OPS if op.SUPPORTS_BMGHK])
 def test_forward_gqa_one_group(opFW):

{xformers-0.0.23.dev644 → xformers-0.0.23.dev648}/xformers/benchmarks/benchmark_attn_decoding.py

@@ -4,12 +4,11 @@
 # LICENSE file in the root directory of this source tree.
 
 
-import itertools
 from typing import Any
 
 import torch
 from torch.utils import benchmark
-from utils import benchmark_main_helper
+from utils import benchmark_main_helper2
 
 import xformers.ops as xops
 
@@ -17,13 +16,6 @@ min_run_time = 0.5
 device = torch.device("cuda")
 
 
-def product_dict(**kwargs):
-    keys = kwargs.keys()
-    vals = kwargs.values()
-    for instance in itertools.product(*vals):
-        yield dict(zip(keys, instance))
-
-
 CASES = [
     dict(B=max(1, 2 ** (16 - i)), Mq=1, Mkv=2**i, Hq=16, Hkv=1, K=128)
     for i in range(8, 18)
@@ -147,13 +139,10 @@ except ImportError:
     pass
 
 
-def attn_decoding(**kwargs):
-    yield from _setup_test(
-        **kwargs,
-        fw=True,
-        cuda_graph=True,
-        functions=BENCHMARKS,
-    )
-
-
-benchmark_main_helper(attn_decoding, CASES, min_run_time=min_run_time)
+benchmark_main_helper2(
+    "attn_decoding",
+    fw=True,
+    cases=CASES,
+    functions=BENCHMARKS,
+    min_run_time=min_run_time,
+)

xformers-0.0.23.dev648/xformers/benchmarks/benchmark_indexing.py

@@ -0,0 +1,191 @@
+# 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 random
+
+import torch
+from utils import DTYPE2STR, benchmark_main_helper2, product_dict
+
+import xformers.ops as xops
+
+min_run_time = 0.2
+device = torch.device("cuda")
+
+CASES_IADD = list(
+    product_dict(
+        shape=[
+            (int(48 * 0.6), 48, 1, 257 * 1536),
+            (int(48 * 0.6), 48, 257, 1536),
+        ],
+        scaling=[False, True],
+        dtype=[torch.half],
+    )
+) + list(
+    product_dict(
+        shape=[
+            # Format: [B_src, B_inp, M, D]
+            (int(192 * 0.6), 192, 50, 1536),
+            (int(48 * 257 * 0.6), 257 * 48, 1, 1536),
+            (int(192 * 50 * 0.6), 192 * 50, 1, 1536),
+            (int(16 * 257 * 0.6), 48 * 257, 1, 1536),
+        ],
+        scaling=[False],
+        dtype=[torch.half],
+    )
+)
+
+CASES_ISELECT = list(
+    product_dict(
+        batches=[((48, 257), (50, 192))],
+        D=[1536],
+        keep_ratio=[0.6],
+        dtype=[torch.half],
+    )
+)
+
+
+class ScaledIndexAddBenchmark:
+    def __init__(self, dtype, scaling: bool, shape, bw: bool) -> None:
+        B_src, B_out, M, D = shape
+        torch.manual_seed(B_out + B_src)
+        dtype_str = DTYPE2STR.get(dtype, dtype)
+        self.sub_label = f"{dtype_str} B_src={B_src}, B_out={B_out}, M={M}, D={D} s={'Y' if scaling else 'N'}"
+        self.label = "scaled_index_add"
+        self.alpha = 0.73
+
+        self.inp = torch.randn(
+            [B_out, M, D], device="cuda", dtype=dtype, requires_grad=bw
+        )
+        self.src = torch.randn(
+            [B_src, M, D], device="cuda", dtype=dtype, requires_grad=bw
+        )
+        self.scaling = (
+            torch.randn([D], device="cuda", dtype=dtype, requires_grad=bw)
+            if scaling
+            else None
+        )
+        self.index = torch.tensor(
+            [i for i in range(self.src.shape[0])], dtype=torch.int64, device="cuda"
+        )
+        self.grad = torch.randn([B_out, M, D], device="cuda", dtype=dtype)
+        self.out = torch.Tensor()
+
+    def fw(self) -> None:
+        self.out = xops.scaled_index_add(
+            input=self.inp.clone(),
+            index=self.index,
+            source=self.src,
+            scaling=self.scaling,
+            alpha=self.alpha,
+        )
+
+    def bw(self):
+        self.inp.grad = None
+        self.src.grad = None
+        if self.scaling is not None:
+            self.scaling.grad = None
+        self.out.backward(self.grad, retain_graph=True)
+
+
+class ScaledIndexAddBenchmarkBaseline(ScaledIndexAddBenchmark):
+    def fw(self) -> None:
+        src_scaled = self.src
+        if self.scaling is not None:
+            src_scaled * self.scaling.unsqueeze(0).unsqueeze(0)
+        self.out = self.inp.index_add(
+            dim=0,
+            source=src_scaled,
+            index=self.index,
+            alpha=self.alpha,
+        )
+
+
+class IndexSelectBenchmark:
+    def __init__(self, dtype, batches, D, keep_ratio, bw: bool) -> None:
+        dtype_str = DTYPE2STR.get(dtype, dtype)
+        self.sub_label = f"{dtype_str} D={D} batches={batches} keep={keep_ratio}"
+        self.label = "index_select"
+
+        indices = []
+        sources = []
+        for (B, seqlen) in batches:
+            index = [i for i in range(B)]
+            random.Random(B).shuffle(index)
+            indices.append(
+                torch.zeros(
+                    index[int(keep_ratio * B)],
+                    dtype=torch.int64,
+                    device="cuda",
+                )
+            )
+            source_i = torch.randn(
+                [B, seqlen * D], dtype=dtype, device="cuda", requires_grad=bw
+            )
+            sources.append(source_i)
+        self.indices, self.sources = indices, sources
+        self.out = torch.Tensor()
+
+    def fw(self) -> None:
+        self.out = xops.index_select_cat(self.sources, self.indices)
+
+    def bw(self):
+        for src in self.sources:
+            src.grad = None
+        self.out.backward(self.out, retain_graph=True)
+
+
+class IndexSelectBenchmarkBaseline(IndexSelectBenchmark):
+    def fw(self) -> None:
+        self.out = torch.cat(
+            [s[i].flatten() for s, i in zip(self.sources, self.indices)], dim=0
+        )
+
+
+benchmark_main_helper2(
+    "scaled_index_add_fw",
+    fw=True,
+    functions={
+        "xformers": ScaledIndexAddBenchmark,
+        "pytorch": ScaledIndexAddBenchmarkBaseline,
+    },
+    cases=CASES_IADD,
+    min_run_time=min_run_time,
+)
+
+benchmark_main_helper2(
+    "scaled_index_add_fwbw",
+    fw=True,
+    bw=True,
+    functions={
+        "xformers": ScaledIndexAddBenchmark,
+        "pytorch": ScaledIndexAddBenchmarkBaseline,
+    },
+    cases=CASES_IADD,
+    min_run_time=min_run_time,
+)
+
+benchmark_main_helper2(
+    "index_select_fw",
+    fw=True,
+    functions={
+        "xformers": IndexSelectBenchmark,
+        "pytorch": IndexSelectBenchmarkBaseline,
+    },
+    cases=CASES_ISELECT,
+    min_run_time=min_run_time,
+)
+
+benchmark_main_helper2(
+    "index_select_fwbw",
+    fw=True,
+    bw=True,
+    functions={
+        "xformers": IndexSelectBenchmark,
+        "pytorch": IndexSelectBenchmarkBaseline,
+    },
+    cases=CASES_ISELECT,
+    min_run_time=min_run_time,
+)

{xformers-0.0.23.dev644 → xformers-0.0.23.dev648}/xformers/benchmarks/utils.py

@@ -8,13 +8,14 @@ import contextlib
 import copy
 import csv
 import glob
+import itertools
 import logging
 import math
 import os
 import tempfile
 from collections import defaultdict, namedtuple
 from dataclasses import replace
-from typing import Any, Dict, Generator, List, Set, Tuple
+from typing import Any, Dict, Generator, Iterator, List, Set, Tuple
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -489,6 +490,10 @@ def benchmark_run_and_compare(
         ):
             loaded = _benchmark_results_from_csv(filename)
             for m, r in loaded:
+                if m.get(META_ALGORITHM) is not None:
+                    suffix_pos = m[META_ALGORITHM].find("@")
+                    if suffix_pos > 0:
+                        m[META_ALGORITHM] = m[META_ALGORITHM][:suffix_pos]
                 if r.task_spec.env == env and SKIP_VANILLA_TASKS_IF_ALREADY_DONE:
                     skip_vanilla_tasks.add(
                         (r.task_spec.sub_label, r.task_spec.num_threads)
@@ -654,7 +659,71 @@ def _fail_if_regressions(
     print(f"  Worse    : {num_worse}")
     if num_unk > 0:
         print(f"  (no ref) : {num_unk}")
+    benchmarks_run = num_better + num_nochange + num_worse
     if num_worse > 1:
         raise RuntimeError("At least one benchmark regressed!")
-    if num_nochange == 0:
+    elif num_unk == benchmarks_run:
         raise RuntimeError("No reference found")
+    elif benchmarks_run == 0:
+        raise RuntimeError("No benchmark was run")
+
+
+def benchmark_main_helper2(
+    name: str,
+    functions,
+    fw: bool = False,
+    bw: bool = False,
+    cuda_graph: bool = True,
+    **kwargs,
+) -> None:
+    assert fw or bw
+
+    def handle_case(**case) -> Iterator[benchmark.Timer]:
+        for k, benchmark_cls in functions.items():
+            benchmark_object = benchmark_cls(**case, bw=bw)
+            label = benchmark_object.label
+            label += "fw" if fw else ""
+            label += "bw" if bw else ""
+
+            def run_one():
+                if fw:
+                    benchmark_object.fw()
+                if bw:
+                    benchmark_object.bw()
+
+            if cuda_graph:
+                run_one()
+                benchmark_object = benchmark_cls(**case, bw=bw)
+                g = torch.cuda.CUDAGraph()
+                with torch.cuda.graph(g):
+                    run_one()
+
+                def run_one():
+                    g.replay()
+
+            yield benchmark.Timer(
+                stmt="fn()",
+                globals={
+                    "fn": run_one,
+                },
+                label=label,
+                description=k,
+                sub_label=benchmark_object.sub_label,
+            )
+
+    handle_case.__name__ = name
+    benchmark_main_helper(handle_case, **kwargs)
+
+
+def product_dict(**kwargs):
+    keys = kwargs.keys()
+    vals = kwargs.values()
+    for instance in itertools.product(*vals):
+        yield dict(zip(keys, instance))
+
+
+DTYPE2STR = {
+    torch.bfloat16: "b16",
+    torch.half: "f16",
+    torch.float32: "f32",
+}

{xformers-0.0.23.dev644 → xformers-0.0.23.dev648}/xformers/ops/fmha/flash.py

@@ -208,6 +208,7 @@ except ImportError:
 
 def _convert_input_format(
     inp: Inputs,
+    supports_mqa: bool,
 ) -> Tuple[Inputs, Optional[torch.Tensor], int, Optional[torch.Tensor], int]:
     assert inp.query.ndim in [4, 5]
     query, key, value = inp.query, inp.key, inp.value
@@ -237,6 +238,8 @@ def _convert_input_format(
         max_seqlen_k = inp.key.shape[1]
 
     if query.ndim == 5:  # QGA
+        assert supports_mqa
+
         # Fold the group/head_in_group dimensions together
         def fold(x):
             # Either the head is replicated
@@ -256,7 +259,7 @@ def _convert_input_format(
         key = fold(key)
         value = fold(value)
     # Optimize for MHA
-    if key.ndim == 4 and key.stride(2) == 0 and value.stride(2) == 0:
+    if key.ndim == 4 and key.stride(2) == 0 and value.stride(2) == 0 and supports_mqa:
         key = key[:, :, :1]
         value = value[:, :, :1]
     # Initially we have `query.shape = [batch, seqlen, head_dim_q]`
@@ -392,7 +395,7 @@ class FwOp(AttentionFwOpBase):
             max_seqlen_q,
             cu_seqlens_k,
             max_seqlen_k,
-        ) = _convert_input_format(inp)
+        ) = _convert_input_format(inp, supports_mqa=True)
         if inp.query.numel() > 0 and inp.key.numel() > 0:
             out, softmax_lse, rng_state = cls.OPERATOR(
                 inp.query,
@@ -510,7 +513,7 @@ class BwOp(AttentionBwOpBase):
             max_seqlen_q,
             cu_seqlens_k,
             max_seqlen_k,
-        ) = _convert_input_format(inp)
+        ) = _convert_input_format(inp, supports_mqa=False)
         assert ctx.lse.is_contiguous()
         ctx_lse = ctx.lse
         assert ctx_lse.shape[2] >= max_seqlen_q

{xformers-0.0.23.dev644 → xformers-0.0.23.dev648}/xformers/ops/indexing.py

@@ -14,6 +14,31 @@ from xformers.ops.triton import (
     scaled_index_add_fwd,
 )
 
+from .common import BaseOperator, register_operator
+
+
+# Keeping these operator registry here so that
+# it's easy to check if they are available
+@register_operator
+class ScaledIndexAddFw(BaseOperator):
+    OPERATOR = scaled_index_add_fwd
+    OPERATOR_CATEGORY = "indexing"
+    NAME = "scaled_index_addF"
+
+
+@register_operator
+class ScaledIndexAddBw(BaseOperator):
+    OPERATOR = scaled_index_add_bwd
+    OPERATOR_CATEGORY = "indexing"
+    NAME = "scaled_index_addB"
+
+
+@register_operator
+class IndexSelect(BaseOperator):
+    OPERATOR = index_select_cat_fwd
+    OPERATOR_CATEGORY = "indexing"
+    NAME = "index_select"
+
 
 class _ScaledIndexAdd(torch.autograd.Function):
     @staticmethod
@@ -77,7 +102,7 @@ class _ScaledIndexAdd(torch.autograd.Function):
 
 
 def scaled_index_add(
-    x: torch.Tensor,  # [B, M, D]
+    input: torch.Tensor,  # [B, M, D]
     index: torch.Tensor,  # [Bi] - int64
     source: torch.Tensor,  # [Bi, M, D]
     scaling: Optional[torch.Tensor] = None,  # [D]
@@ -88,20 +113,20 @@ def scaled_index_add(
 
     Indices in ``index`` are assumed to be unique
 
-    The max index in ``index`` is assumed to be less than the size of dim0 of input ``x``.
+    The max index in ``index`` is assumed to be less than the size of dim0 of ``input``.
 
     :Note:
 
-        The FW pass is done in-place (input ``x`` is modified)
+        The FW pass is done in-place (``input`` is modified)
 
     :Equivalent pytorch code:
 
     .. code-block:: python
 
-        return torch.index_add(inp, dim=0, source=scaling * src, index=indices, alpha=alpha)
+        return torch.index_add(input, dim=0, source=scaling * src, index=indices, alpha=alpha)
     """
 
-    return _ScaledIndexAdd.apply(x, index, source, scaling, alpha)
+    return _ScaledIndexAdd.apply(input, index, source, scaling, alpha)
 
 
 class _IndexSelectCat(torch.autograd.Function):

xformers-0.0.23.dev648/xformers/version.py

@@ -0,0 +1,2 @@
+# noqa: C801
+__version__ = "0.0.23.dev648"

{xformers-0.0.23.dev644 → xformers-0.0.23.dev648/xformers.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xformers
-Version: 0.0.23.dev644
+Version: 0.0.23.dev648
 Summary: XFormers: A collection of composable Transformer building blocks.
 Home-page: https://facebookresearch.github.io/xformers/
 Author: Facebook AI Research