xformers 0.0.24.dev735__tar.gz → 0.0.25.dev740__tar.gz

{xformers-0.0.24.dev735/xformers.egg-info → xformers-0.0.25.dev740}/PKG-INFO

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

{xformers-0.0.24.dev735 → xformers-0.0.25.dev740}/README.md

@@ -28,13 +28,13 @@ xFormers is:
 
 ## Installing xFormers
 
-* **(RECOMMENDED, linux) Install latest stable with conda**: Requires [PyTorch 2.1.2 installed with conda](https://pytorch.org/get-started/locally/)
+* **(RECOMMENDED, linux) Install latest stable with conda**: Requires [PyTorch 2.2.0 installed with conda](https://pytorch.org/get-started/locally/)
 
 ```bash
 conda install xformers -c xformers
 ```
 
-* **(RECOMMENDED, linux & win) Install latest stable with pip**: Requires [PyTorch 2.1.2](https://pytorch.org/get-started/locally/)
+* **(RECOMMENDED, linux & win) Install latest stable with pip**: Requires [PyTorch 2.2.0](https://pytorch.org/get-started/locally/)
 
 ```bash
 # cuda 11.8 version

xformers-0.0.25.dev740/version.txt

@@ -0,0 +1 @@
+0.0.25.dev740

xformers-0.0.25.dev740/xformers/benchmarks/benchmark_sp24.py

@@ -0,0 +1,178 @@
+# 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.
+
+
+from typing import Tuple
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from utils import DTYPE2STR, benchmark_main_helper2, product_dict
+
+import xformers.ops as xops
+
+min_run_time = 0.5
+device = torch.device("cuda")
+
+CASES = list(
+    product_dict(
+        B_in_hidden_out_ft=[
+            (2048 * 8, 2048, 2048 * 3, 2048),
+            (2048, 5120, 5120 * 3, 5120),  # 13b
+            (1024, 8192, 8192 * 3, 8192),  # 30b
+            (2048, 8192, 8192 * 3, 8192),  # 30b
+            (2048 * 2, 8192, 8192 * 3, 8192),  # 30b
+            # DINO ViT-L: lg + sm crops (patch16)
+            (64 * 2 * (14 * 14 + 1) + 64 * 8 * (6 * 6 + 1), 1024, 1024 * 4, 1024),
+            # DINO ViT-g: lg + sm crops (patch16)
+            (
+                12 * 2 * (16 * 16 + 1 + 11) + 12 * 8 * (7 * 7 + 1 + 11),
+                1536,
+                1536 * 4,
+                1536,
+            ),
+        ],
+        dtype=[torch.half],
+        bias=[False],
+    )
+)
+
+
+class Mlp(nn.Module):
+    LINEAR_CLS = nn.Linear
+
+    def __init__(
+        self, B_in_hidden_out_ft: Tuple[int, int, int, int], dtype, bias: bool, bw: bool
+    ) -> None:
+        B, in_ft, hid_ft, out_ft = B_in_hidden_out_ft
+        super().__init__()
+        self.label = "mlp"
+        self.sub_label = (
+            f"{DTYPE2STR[dtype]} ({B},{in_ft},{hid_ft},{out_ft}){' b' if bias else ''}"
+        )
+        self.fc1 = self.LINEAR_CLS(in_ft, hid_ft, bias=bias)
+        self.act = nn.GELU()
+        self.fc2 = self.LINEAR_CLS(hid_ft, out_ft, bias=bias)
+        self.grad = torch.randn([B, out_ft], device="cuda", dtype=dtype)
+        self.input = torch.randn(
+            [B, in_ft], device="cuda", dtype=dtype, requires_grad=True
+        )
+        self.out = self.input
+        self.to("cuda").to(dtype)
+
+    def fw(self):
+        x = self.input
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.fc2(x)
+        self.out = x
+
+    def bw(self):
+        self.out.backward(self.grad, retain_graph=True)
+
+
+class MlpDenseMask(Mlp):
+    def fw(self):
+        x = self.input
+        x = self.fc1(x)
+
+        mask = torch.ops.xformers.sparse24_largest_mask_2d(x)
+        x = mask * x
+
+        x = self.act(x)
+        x = self.fc2(x)
+        self.out = x
+
+
+class MlpAct24(Mlp):
+    def fw(self):
+        x = self.input
+        x = self.fc1(x)
+
+        x = xops.sparsify24(x)
+
+        x = self.act(x)
+        x = self.fc2(x)
+        self.out = x
+
+
+class LinearW24(torch.nn.Linear):
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        w_sparse = xops.sparsify24(
+            self.weight,
+            gradient="24dense",
+            backend="cusparselt",
+        )
+        return F.linear(input, w_sparse, self.bias)
+
+
+class MlpW24(Mlp):
+    LINEAR_CLS = LinearW24
+
+
+class MicrobenchmarkBase:
+    def __init__(
+        self, B_in_hidden_out_ft: Tuple[int, int, int, int], dtype, bias: bool, bw: bool
+    ) -> None:
+        B, in_ft, hid_ft, out_ft = B_in_hidden_out_ft
+        super().__init__()
+        self.label = "mlp"
+        self.sub_label = (
+            f"{DTYPE2STR[dtype]} ({B},{in_ft},{hid_ft},{out_ft}){' b' if bias else ''}"
+        )
+        self.input = torch.randn(
+            [B, in_ft], device="cuda", dtype=dtype, requires_grad=True
+        )
+        self.input_colMajor = self.input.t().contiguous().t()
+        self.input_sp = xops.sparsify24(self.input)
+
+    def bw(self) -> None:
+        return None
+
+
+class MicrobenchmarkSparsify24(MicrobenchmarkBase):
+    def fw(self) -> torch.Tensor:
+        xops.sparsify24(self.input)
+        return self.input
+
+
+class MicrobenchmarkSp24ApplyDense(MicrobenchmarkBase):
+    def fw(self) -> torch.Tensor:
+        xops.sparsify24_like(self.input, pattern=self.input_sp, out_dense=True)
+        return self.input
+
+
+class MicrobenchmarkSp24ApplyDenseT(MicrobenchmarkBase):
+    def fw(self) -> torch.Tensor:
+        xops.sparsify24_like(self.input_colMajor, pattern=self.input_sp, out_dense=True)
+        return self.input
+
+
+class MicrobenchmarkInputClone(MicrobenchmarkBase):
+    def fw(self) -> torch.Tensor:
+        self.input.clone()
+        return self.input
+
+
+functions = {
+    "act24": MlpAct24,
+    "dense": Mlp,
+    "w24": MlpW24,
+    "s24_inp_sparsify24": MicrobenchmarkSparsify24,
+    "s24_inp_apply_dense": MicrobenchmarkSp24ApplyDense,
+    "s24_inp_apply_dense_t": MicrobenchmarkSp24ApplyDenseT,
+    "s24_inp_clone": MicrobenchmarkInputClone,
+}
+benchmark_main_helper2(
+    "sp24_fw", fw=True, cases=CASES, functions=functions, min_run_time=min_run_time
+)
+benchmark_main_helper2(
+    "sp24_fwbw",
+    fw=True,
+    bw=True,
+    cases=CASES,
+    functions=functions,
+    min_run_time=min_run_time,
+)

xformers-0.0.25.dev740/xformers/csrc/sparse24/compute_sparse_tile.h

@@ -0,0 +1,187 @@
+#pragma once
+
+#include <cutlass/bfloat16.h>
+#include <cutlass/half.h>
+
+#include "sparse24_pack.h"
+#include "static_sort.h"
+
+// Given 4x4 values, computes the selected indices that will remain after 2:4
+// sparsification, as a bitmask.
+// NOTE: Algorithms might select LESS than 8 values in total in some cases.
+
+namespace platform {
+template <>
+struct numeric_limits<cutlass::bfloat16_t> {
+  CUTLASS_HOST_DEVICE
+  static cutlass::bfloat16_t infinity() {
+    return cutlass::bfloat16_t::bitcast(0x7f80);
+  }
+};
+} // namespace platform
+namespace xformers {
+namespace sp24 {
+
+template <typename Element, typename Pointwise>
+struct TileValueOrderedT {
+  union {
+    struct {
+      Element value;
+      uint2b_t col;
+      uint2b_t row;
+    } parts;
+    uint32_t raw;
+  };
+  CUTLASS_DEVICE bool operator<(
+      TileValueOrderedT<Element, Pointwise> const& other) const {
+    return Pointwise::apply(parts.value) < Pointwise::apply(other.parts.value);
+  }
+  CUTLASS_DEVICE TileValueOrderedT() {}
+};
+
+// Operations that we can apply to rank the values
+struct IdentityOp {
+  template <typename T>
+  static T CUTLASS_HOST_DEVICE apply(T const& x) {
+    return x;
+  }
+};
+// Can be applied to rank based on absolute value
+struct AbsOp {
+  template <typename T>
+  static T CUTLASS_HOST_DEVICE apply(T const& x) {
+    return cutlass::abs(x);
+  }
+};
+
+// Given 4x4 values, computes the selected indices that will remain after 2:4
+// sparsification, as a bitmask. We have 2 constraints:
+// (1) At most 2 values per line
+// (2) At most 2 values per column
+// This means we can select at most 8 values in total.
+// ALGO: We use a greedy algorithm, where we take values in the 4x4
+// tile in descending order. If a value fits (because the line/col is not
+// already full), we select it. Then we move on to the next one.
+// NOTE: This algorithm might select LESS than 8 values in total in some cases.
+// NOTE (2): RF are not indexable, so we shouldn't rely on indexing
+//   values at any point, otherwise they will be stored in local memory.
+template <typename Op = IdentityOp>
+struct LargestValuesGreedy {
+  template <typename T>
+  static CUTLASS_DEVICE T outOfBoundsFillValue() {
+    return -platform::numeric_limits<T>::infinity();
+  }
+
+  template <typename Tile4x4Accessor>
+  CUTLASS_DEVICE Indices4x4 operator()(Tile4x4Accessor values) {
+    using TileValueOrdered =
+        TileValueOrderedT<typename Tile4x4Accessor::Element, Op>;
+    using TileValuesFragment = cutlass::Array<TileValueOrdered, 4 * 4>;
+    Indices4x4 indices;
+    TileValuesFragment values_ordered;
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < 4; ++i) {
+      CUTLASS_PRAGMA_UNROLL
+      for (int j = 0; j < 4; ++j) {
+        TileValueOrdered& v = values_ordered[i * 4 + j];
+        v.parts.value = values.at(i, j).get();
+        v.parts.col = j;
+        v.parts.row = i;
+      }
+    }
+    // Use a sorting network (aka without branches) to avoid
+    // warp divergence
+    StaticSort<TileValuesFragment::kElements> sorter;
+    sorter(values_ordered);
+
+    // bitmask to store how many we have selected on a given row/col
+    // 0 selected: (numPerRow >> 2*row) = 00 (0)
+    // 1 selected: (numPerRow >> 2*row) = 01 (1)
+    // 2 selected: (numPerRow >> 2*row) = 11 (3)
+    uint32_t numPerRow = 0;
+    uint32_t numPerCol = 0;
+    indices = 0;
+
+    // Take as many as we can, starting with the largest values
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = values_ordered.size() - 1; i >= 0; i--) {
+      auto& e = values_ordered[i];
+
+      uint32_t rcount = uint2b_t(numPerRow >> 2 * e.parts.row);
+      uint32_t ccount = uint2b_t(numPerCol >> 2 * e.parts.col);
+      // NOTE: This is more efficient (yet equivalent) to:
+      // `rcount != 3 && ccount != 3`
+      bool selected = (rcount + ccount) <= 2;
+      indices |= selected << (e.parts.col + 4 * e.parts.row);
+
+      numPerRow |= (rcount + selected) << 2 * e.parts.row;
+      numPerCol |= (ccount + selected) << 2 * e.parts.col;
+    }
+    return indices;
+  }
+};
+
+// We consider each rows independantly in order
+// This is to ensure that a row's sparsity pattern is only determined
+// by its values and the rows before (but never the rows after)
+// This enforces causality strictly
+template <typename Op = IdentityOp>
+struct Causal1122 {
+  template <typename T>
+  static CUTLASS_DEVICE T outOfBoundsFillValue() {
+    return -platform::numeric_limits<T>::infinity();
+  }
+
+  template <typename Tile4x4Accessor>
+  CUTLASS_DEVICE Indices4x4 operator()(Tile4x4Accessor values) {
+    static constexpr int kMaxValuesPerRow[] = {1, 1, 2, 2};
+    using TileValueOrdered =
+        TileValueOrderedT<typename Tile4x4Accessor::Element, Op>;
+    using TileValuesFragment = cutlass::Array<TileValueOrdered, 4>;
+    Indices4x4 indices = 0;
+
+    uint32_t numPerCol = 0; // <- see doc in `LargestValuesGreedy`
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int row = 0; row < 4; ++row) {
+      int row_count = 0;
+      TileValuesFragment values_ordered;
+      CUTLASS_PRAGMA_UNROLL
+      for (int col = 0; col < 4; ++col) {
+        TileValueOrdered& v = values_ordered[col];
+        v.parts.value = values.at(row, col).get();
+        v.parts.col = col;
+      }
+      // Use a sorting network (aka without branches) to avoid
+      // warp divergence
+      StaticSort<TileValuesFragment::kElements> sorter;
+      sorter(values_ordered);
+
+      // Take as many as we can, starting with the largest values
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = values_ordered.size() - 1; i >= 0; i--) {
+        auto& e = values_ordered[i];
+
+        uint32_t ccount = uint2b_t(numPerCol >> 2 * e.parts.col);
+        bool selected = ccount != 3 && (row_count < kMaxValuesPerRow[row]);
+        indices |= selected << (e.parts.col + 4 * row);
+        numPerCol |= (ccount + selected) << 2 * e.parts.col;
+        row_count += selected;
+      }
+    }
+    return indices;
+  }
+};
+
+template <typename T>
+void named_algorithms(T callback) {
+  callback(LargestValuesGreedy<IdentityOp>(), "largest_values_greedy");
+  callback(Causal1122<IdentityOp>(), "causal1122");
+  callback(LargestValuesGreedy<AbsOp>(), "largest_abs_values_greedy");
+
+  // default one
+  callback(LargestValuesGreedy<IdentityOp>(), "");
+}
+
+} // namespace sp24
+} // namespace xformers