xformers 0.0.24.dev720__tar.gz → 0.0.24.dev725__tar.gz

{xformers-0.0.24.dev720/xformers.egg-info → xformers-0.0.24.dev725}/PKG-INFO

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

{xformers-0.0.24.dev720 → xformers-0.0.24.dev725}/requirements.txt

@@ -1,4 +1,4 @@
 # Example requirement, can be anything that pip knows
 # install with `pip install -r requirements.txt`, and make sure that CI does the same
-torch >= 1.12
+torch >= 2.1
 numpy

{xformers-0.0.24.dev720 → xformers-0.0.24.dev725}/setup.py

@@ -31,9 +31,15 @@ from torch.utils.cpp_extension import (
 this_dir = os.path.dirname(__file__)
 
 
-def get_extra_nvcc_flags_for_build_type() -> List[str]:
+def get_extra_nvcc_flags_for_build_type(cuda_version: int) -> List[str]:
     build_type = os.environ.get("XFORMERS_BUILD_TYPE", "RelWithDebInfo").lower()
     if build_type == "relwithdebinfo":
+        if cuda_version >= 1201 and cuda_version < 1202:
+            print(
+                "Looks like we are using CUDA 12.1 which segfaults when provided with"
+                " the -generate-line-info flag. Disabling it."
+            )
+            return []
         return ["--generate-line-info"]
     elif build_type == "release":
         return []
@@ -201,7 +207,7 @@ def get_flash_attention_extensions(cuda_version: int, extra_compile_args):
                 ]
                 + nvcc_archs_flags
                 + nvcc_windows_flags
-                + get_extra_nvcc_flags_for_build_type(),
+                + get_extra_nvcc_flags_for_build_type(cuda_version),
             },
             include_dirs=[
                 p.absolute()
@@ -252,6 +258,7 @@ def get_extensions():
         or os.getenv("FORCE_CUDA", "0") == "1"
         or os.getenv("TORCH_CUDA_ARCH_LIST", "") != ""
     ):
+        cuda_version = get_cuda_version(CUDA_HOME)
         extension = CUDAExtension
         sources += source_cuda
         include_dirs += [sputnik_dir, cutlass_dir, cutlass_examples_dir]
@@ -263,11 +270,10 @@ def get_extensions():
             "--extended-lambda",
             "-D_ENABLE_EXTENDED_ALIGNED_STORAGE",
             "-std=c++17",
-        ] + get_extra_nvcc_flags_for_build_type()
+        ] + get_extra_nvcc_flags_for_build_type(cuda_version)
         if os.getenv("XFORMERS_ENABLE_DEBUG_ASSERTIONS", "0") != "1":
             nvcc_flags.append("-DNDEBUG")
         nvcc_flags += shlex.split(os.getenv("NVCC_FLAGS", ""))
-        cuda_version = get_cuda_version(CUDA_HOME)
         if cuda_version >= 1102:
             nvcc_flags += [
                 "--threads",

xformers-0.0.24.dev725/version.txt

@@ -0,0 +1 @@
+0.0.24.dev725

{xformers-0.0.24.dev720 → xformers-0.0.24.dev725}/xformers/__init__.py

@@ -51,9 +51,10 @@ def _is_triton_available():
         from xformers.triton.softmax import softmax as triton_softmax  # noqa
 
         return True
-    except (ImportError, AttributeError) as e:
+    except (ImportError, AttributeError):
         logger.warning(
-            f"A matching Triton is not available, some optimizations will not be enabled.\nError caught was: {e}"
+            "A matching Triton is not available, some optimizations will not be enabled",
+            exc_info=True,
         )
         return False
 

{xformers-0.0.24.dev720 → xformers-0.0.24.dev725}/xformers/checkpoint.py

@@ -4,8 +4,8 @@
 # LICENSE file in the root directory of this source tree.
 
 
+import functools
 import time
-import weakref
 from collections import defaultdict
 from contextlib import nullcontext
 from copy import deepcopy
@@ -20,7 +20,6 @@ from torch.testing._internal.composite_compliance import (
 )
 from torch.utils._python_dispatch import TorchDispatchMode
 from torch.utils._pytree import tree_map
-from torch.utils.checkpoint import get_device_states, set_device_states
 
 _scipy_is_available = False
 try:
@@ -35,13 +34,28 @@ try:
     from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
         ActivationWrapper,
     )
+    from torch.utils.checkpoint import (  # type: ignore
+        _CachedTorchDispatchMode,
+        _CachingTorchDispatchMode,
+        _ignored_ops,
+    )
 except ImportError:
     ActivationWrapper = torch.nn.Module  # type: ignore
 
+    class _NotAvailable:
+        def __init__(self, *args, **kwargs):
+            raise RuntimeError("Need PyTorch >= 2.2")
+
+    _CachedTorchDispatchMode = _NotAvailable  # type: ignore
+    _CachingTorchDispatchMode = _NotAvailable  # type: ignore
+    _ignored_ops = set()  # type: ignore
+
 
-OPS_TO_ALWAYS_SKIP = {
-    "aten.detach.default",
+_additional_ignored_ops = {
+    torch.ops.aten.lift_fresh.default,
+    torch.ops.profiler._record_function_exit._RecordFunction,
 }
+OPS_TO_ALWAYS_SKIP = _ignored_ops | _additional_ignored_ops
 
 
 @dataclass
@@ -62,41 +76,6 @@ def _detach(x):
     return x
 
 
-class CachingTorchDispatchMode(TorchDispatchMode):
-    def __init__(self, policy_fn, storage):
-        self.policy_fn = policy_fn
-        self.storage = storage
-
-    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
-        if kwargs is None:
-            kwargs = {}
-        if self.policy_fn(func, *args, **kwargs):
-            out = func(*args, **kwargs)
-            out_detached = tree_map(_detach, out)
-            self.storage[func].append(out_detached)
-            return out
-        return func(*args, **kwargs)
-
-
-class CachedTorchDispatchMode(TorchDispatchMode):
-    def __init__(self, policy_fn, storage):
-        self.policy_fn = policy_fn
-        self.storage = storage
-
-    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
-        if kwargs is None:
-            kwargs = {}
-        if self.policy_fn(func, *args, **kwargs):
-            # this is a basic guard if there are additional ops
-            # that were present in the second forward. An example
-            # is for ops like `detach`, which can be added if our
-            # policy is too loose
-            if self.storage[func]:
-                out = self.storage[func].pop(0)
-                return out
-        return func(*args, **kwargs)
-
-
 def _get_default_policy(allow_list=None):
     _default_allow_list = [
         "xformers.efficient_attention_forward_cutlass.default",
@@ -135,20 +114,19 @@ def list_operators(function, *args, **kwargs):
     return verbose_mode.operators
 
 
-def checkpoint(
-    function, *args, preserve_rng_state=True, policy_fn=None, **kwargs
-) -> Any:
-    """Checkpointining with custom policy function for selectively deciding
-    what to store and what to recompute
+class CachedTorchDispatchMode(_CachedTorchDispatchMode):
+    def pop_from_storage(self, func, args, kwargs):
+        # the autograd engine might add spurious views. This is a basic
+        # guard and should be improved
+        if self.storage[func]:
+            return self.storage[func].pop(0)
+        return func(*args, **kwargs)
+
+
+def selective_checkpoint_context_fn(policy_fn=None):
+    """An activation checkpoint context_fn for selectively deciding what to
+    store and what to recompute. Accepts a custom policy.
     Args:
-        function: describes what to run in the forward pass of the model or
-            part of the model. It should also know how to handle the inputs
-            passed as the tuple. For example, in LSTM, if user passes
-            ``(activation, hidden)``, :attr:`function` should correctly use the
-            first input as ``activation`` and the second input as ``hidden``
-        preserve_rng_state(bool, optional):  Omit stashing and restoring
-            the RNG state during each checkpoint.
-            Default: ``True``
         policy_fn(Union[List[Op], callable]): policy for deciding what to
             store (instead of recompute). If it's a function, it should
             be of form (func, *args, **kwargs) -> bool which indicates
@@ -156,40 +134,7 @@ def checkpoint(
             Additionally, a list[Op] is also supported for easier cases.
             The op should be in the format `torch.ops.***`, where the `***`
             names of operators can be obtained with `list_operators`.
-        *args: Arguments to pass in to the given ``function``.
-        **kwargs: Keyword arguments to pass into the given ``function``.
     """
-
-    # Requires PyTorch 1.13 at least
-    from torch.utils.checkpoint import _get_autocast_kwargs
-
-    # Accommodates the (remote) possibility that autocast is enabled for cpu AND gpu.
-    gpu_autocast_kwargs, cpu_autocast_kwargs = _get_autocast_kwargs()
-
-    if preserve_rng_state:
-        fwd_cpu_state = torch.get_rng_state()
-        # Don't eagerly initialize the cuda context by accident.
-        # (If the user intends that the context is initialized later, within their
-        # run_function, we SHOULD actually stash the cuda state here.  Unfortunately,
-        # we have no way to anticipate this will happen before we run the function.
-        # If they do so, we raise an error.)
-        had_cuda_in_fwd = False
-        if torch.cuda._initialized:
-            had_cuda_in_fwd = True
-            fwd_gpu_devices, fwd_gpu_states = get_device_states(*args)
-
-    # Custom class to be able to take weak references
-    class Holder:
-        pass
-
-    # The Holder object for each of the saved object is saved directly on the
-    # SavedVariable and is cleared when reset_data() is called on it. We MUST make
-    # sure that this is the only object having an owning reference to ensure that
-    # the Tensor stored in storage is deleted as soon as the corresponding SavedVariable
-    # data is cleared.
-    storage: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary()
-    weak_holder_list = []
-
     if policy_fn is None:
         policy_fn = _get_default_policy()
     elif isinstance(policy_fn, list):
@@ -201,81 +146,45 @@ def checkpoint(
     # assumption: grad_mode doesn't change inside function
     caching_mode: ContextManager[None]
     if torch.is_grad_enabled():
-        caching_mode = CachingTorchDispatchMode(deepcopy(policy_fn), temp_storage)
+        caching_mode = _CachingTorchDispatchMode(deepcopy(policy_fn), temp_storage)
     else:
         caching_mode = nullcontext()
     cached_mode = CachedTorchDispatchMode(deepcopy(policy_fn), temp_storage)
 
-    def pack(x):
-        # TODO(varal7): Instead of returning abstract object, we can return things metadata (such as
-        # size, device, ...) to catch certain cases of undeterministic behavior of the forward
-        res = Holder()
-        weak_holder_list.append(weakref.ref(res))
-        return res
-
-    def unpack(x):
-        unpack_counter = 0
-        if len(storage) == 0:
-
-            def inner_pack(inner):
-                nonlocal unpack_counter
-                unpack_counter += 1
-                # If the holder went out of scope, the SavedVariable is dead and so
-                # the value will never be read from the storage. Skip filling it.
-                if weak_holder_list[unpack_counter - 1]() is None:
-                    return
-                # Use detach here to ensure we don't keep the temporary autograd
-                # graph created during the second forward
-                storage[weak_holder_list[unpack_counter - 1]()] = inner.detach()
-                return
-
-            def inner_unpack(packed):
-                raise RuntimeError(
-                    "You are calling backwards on a tensor that is never exposed. Please open an issue."
-                )
-
-            # Stash the surrounding rng state, and mimic the state that was
-            # present at this time during forward.  Restore the surrounding state
-            # when we're done.
-            rng_devices = []
-            if preserve_rng_state and had_cuda_in_fwd:
-                rng_devices = fwd_gpu_devices
-            with torch.random.fork_rng(devices=rng_devices, enabled=preserve_rng_state):
-                if preserve_rng_state:
-                    torch.set_rng_state(fwd_cpu_state)
-                    if had_cuda_in_fwd:
-                        set_device_states(fwd_gpu_devices, fwd_gpu_states)
-
-                with cached_mode, torch.enable_grad(), torch.cuda.amp.autocast(
-                    **gpu_autocast_kwargs
-                ), torch.cpu.amp.autocast(
-                    **cpu_autocast_kwargs
-                ), torch.autograd.graph.saved_tensors_hooks(
-                    inner_pack, inner_unpack
-                ):
-                    _unused = function(*args, **kwargs)  # noqa: F841
-
-        if x not in storage:
-            raise RuntimeError(
-                "Attempt to retrieve a tensor saved by autograd multiple times without checkpoint"
-                " recomputation being triggered in between, this is not currently supported. Please"
-                " open an issue with details on your use case so that we can prioritize adding this."
-            )
+    return caching_mode, cached_mode
 
-        return storage[x]
-
-    with caching_mode, torch.autograd.graph.saved_tensors_hooks(pack, unpack):
-        output = function(*args, **kwargs)
-        if torch.cuda._initialized and preserve_rng_state and not had_cuda_in_fwd:
-            # Cuda was not initialized before running the forward, so we didn't
-            # stash the CUDA state.
-            raise RuntimeError(
-                "PyTorch's CUDA state was initialized in the forward pass "
-                "of a Checkpoint, which is not allowed. Please open an issue "
-                "if you need this feature."
-            )
 
-    return output
+def checkpoint(
+    function, *args, preserve_rng_state=True, policy_fn=None, **kwargs
+) -> Any:
+    """Wrapper around torch.utils.checkpoint that accepts a custom policy
+    function for selectively deciding what to store and what to recompute
+    Args:
+        function: describes what to run in the forward pass of the model or
+            part of the model. It should also know how to handle the inputs
+            passed as the tuple. For example, in LSTM, if user passes
+            ``(activation, hidden)``, :attr:`function` should correctly use the
+            first input as ``activation`` and the second input as ``hidden``
+        preserve_rng_state(bool, optional):  Omit stashing and restoring
+            the RNG state during each checkpoint.
+            Default: ``True``
+        policy_fn(Union[List[Op], callable]): policy for deciding what to
+            store (instead of recompute). If it's a function, it should
+            be of form (func, *args, **kwargs) -> bool which indicates
+            if func outputs with *args and **kwargs should be stored or not.
+            Additionally, a list[Op] is also supported for easier cases.
+            The op should be in the format `torch.ops.***`, where the `***`
+            names of operators can be obtained with `list_operators`.
+        *args: Arguments to pass in to the given ``function``.
+        **kwargs: Keyword arguments to pass into the given ``function``.
+    """
+    return torch.utils.checkpoint.checkpoint(
+        function,
+        *args,
+        use_reentrant=False,
+        context_fn=functools.partial(selective_checkpoint_context_fn, policy_fn),
+        **kwargs,
+    )
 
 
 class ProfileOperatorsTorchDispatchMode(TorchDispatchMode):
@@ -339,7 +248,7 @@ class ProfileOperatorsTorchDispatchMode(TorchDispatchMode):
 
         self.data.append(
             ProfileMetadata(
-                str(func),
+                func,
                 time_taken,
                 mem2 - mem1,
                 curr_idx,
@@ -515,9 +424,9 @@ class _OptimalPolicy:
         self.counter = 0
         self.optim_output = optim_output.tolist()
 
-    def __call__(self, func, *args, **kwargs) -> bool:
+    def __call__(self, mode, func, *args, **kwargs) -> bool:
         # returning False means recompute, True means store in memory
-        if str(func) in OPS_TO_ALWAYS_SKIP:
+        if func in OPS_TO_ALWAYS_SKIP:
             return False
         count = self.counter
         self.counter += 1

{xformers-0.0.24.dev720 → xformers-0.0.24.dev725}/xformers/csrc/swiglu/cuda/dual_gemm_silu_identity_mul.cu

@@ -16,6 +16,28 @@
 #include <45_dual_gemm/thread/left_silu_and_mul.h>
 
 namespace {
+
+std::tuple<at::Tensor, at::Tensor, at::Tensor> dual_gemm_silu_identity_mul_META(
+    const at::Tensor& x,
+    const at::Tensor& w0,
+    const c10::optional<at::Tensor>& b0,
+    const at::Tensor& w1,
+    const c10::optional<at::Tensor>& b1) {
+  TORCH_CHECK(x.sym_stride(-1) == 1);
+  TORCH_CHECK(w0.sym_stride(-1) == 1);
+  TORCH_CHECK(w1.sym_stride(-1) == 1);
+
+  at::SymInt B = x.sym_size(0);
+  at::SymInt I = x.sym_size(1);
+  at::SymInt H = w0.sym_size(0);
+
+  at::Tensor d0 = at::empty_symint({B, H}, x.options());
+  at::Tensor d1 = at::empty_symint({B, H}, x.options());
+  at::Tensor d2 = at::empty_symint({B, H}, x.options());
+
+  return std::make_tuple(d0, d1, d2);
+}
+
 template <typename scalar_t>
 std::tuple<at::Tensor, at::Tensor, at::Tensor> dual_gemm_silu_identity_mul_(
     const at::Tensor& x,
@@ -23,12 +45,15 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> dual_gemm_silu_identity_mul_(
     const c10::optional<at::Tensor>& b0,
     const at::Tensor& w1,
     const c10::optional<at::Tensor>& b1) {
+  TORCH_CHECK(x.dim() == 2);
+  TORCH_CHECK(w0.dim() == 2);
+  TORCH_CHECK(w1.dim() == 2);
+
   TORCH_CHECK(x.stride(-1) == 1);
   TORCH_CHECK(w0.stride(-1) == 1);
   TORCH_CHECK(w1.stride(-1) == 1);
 
   at::cuda::CUDAGuard device_guard(x.device());
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
   int64_t B = x.size(0);
   int64_t I = x.size(1);
@@ -38,6 +63,8 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> dual_gemm_silu_identity_mul_(
   at::Tensor d1 = at::empty({B, H}, x.options());
   at::Tensor d2 = at::empty({B, H}, x.options());
 
+  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
   // templati-ze the cutlass kernel
   cutlass::gemm::GemmCoord problem_size(B, H, I);
 
@@ -149,6 +176,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> dual_gemm_silu_identity_mul_(
       typename DualGemm::EpilogueOutputOp1::Params{alpha1, beta1},
       typename DualGemm::EpilogueOutputOp2::Params{},
       split_k_slices};
+
   DualGemm dual_gemm;
   at::Tensor workspace = at::empty(
       {int64_t(dual_gemm.get_workspace_size(arguments))},
@@ -158,10 +186,12 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> dual_gemm_silu_identity_mul_(
       status == cutlass::Status::kSuccess,
       "`dual_gemm_silu_identity_mul` does not support this input: ",
       cutlass::cutlassGetStatusString(status));
+
   status = dual_gemm.initialize(arguments, (uint8_t*)workspace.data_ptr());
   TORCH_CHECK(status == cutlass::Status::kSuccess, "kernel initialize failed");
   status = dual_gemm(stream);
   TORCH_CHECK(status == cutlass::Status::kSuccess, "kernel run failed");
+
   return std::make_tuple(d0, d1, d2);
 }
 
@@ -177,7 +207,6 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> dual_gemm_silu_identity_mul(
   TORCH_CHECK(w1.dim() == 2);
 
 #define FWD_PARAMS x, w0, b0, w1, b1
-
   if (x.scalar_type() == at::ScalarType::Half) {
     return dual_gemm_silu_identity_mul_<cutlass::half_t>(FWD_PARAMS);
   } else {
@@ -212,6 +241,12 @@ TORCH_LIBRARY_IMPL(xformers, CUDA, m) {
       TORCH_FN(dual_gemm_silu_identity_mul));
 }
 
+TORCH_LIBRARY_IMPL(xformers, Meta, m) {
+  m.impl(
+      TORCH_SELECTIVE_NAME("xformers::dual_gemm_silu_identity_mul"),
+      TORCH_FN(dual_gemm_silu_identity_mul_META));
+}
+
 TORCH_LIBRARY_IMPL(xformers, Autocast, m) {
   m.impl(
       TORCH_SELECTIVE_NAME("xformers::dual_gemm_silu_identity_mul"),

{xformers-0.0.24.dev720 → xformers-0.0.24.dev725}/xformers/csrc/swiglu/cuda/gemm_fused_operand_sum.cu

@@ -202,6 +202,7 @@ void gemm_fused_operand_sum_(
   TORCH_CHECK(status == cutlass::Status::kSuccess, "kernel run failed");
 }
 
+template <bool kIsMeta = false>
 std::tuple<at::Tensor, at::Tensor> gemm_fused_operand_sum(
     const at::Tensor& a,
     const at::Tensor& b,
@@ -211,24 +212,27 @@ std::tuple<at::Tensor, at::Tensor> gemm_fused_operand_sum(
   TORCH_CHECK(a.dim() == 2);
   TORCH_CHECK(b.dim() == 2);
   TORCH_CHECK(out_mm.dim() == 2);
-  TORCH_CHECK(out_mm.size(0) == a.size(0));
-  TORCH_CHECK(out_mm.size(1) == b.size(1));
+  TORCH_CHECK(out_mm.sym_size(0) == a.sym_size(0));
+  TORCH_CHECK(out_mm.sym_size(1) == b.sym_size(1));
   TORCH_CHECK(out_sum.dim() == 1);
 
 #define FWD_PARAMS a, b, out_mm, out_sum
 
-  if (a.scalar_type() == at::ScalarType::Half) {
-    TORCH_CHECK(b.scalar_type() == at::ScalarType::Half);
-    TORCH_CHECK(out_mm.scalar_type() == at::ScalarType::Half);
-    TORCH_CHECK(out_sum.scalar_type() == at::ScalarType::Half);
-    gemm_fused_operand_sum_<cutlass::half_t>(FWD_PARAMS);
-  } else {
-    TORCH_CHECK(
-        a.scalar_type() == at::ScalarType::BFloat16, "Only supports bf16/f16");
-    TORCH_CHECK(b.scalar_type() == at::ScalarType::BFloat16);
-    TORCH_CHECK(out_mm.scalar_type() == at::ScalarType::BFloat16);
-    TORCH_CHECK(out_sum.scalar_type() == at::ScalarType::BFloat16);
-    gemm_fused_operand_sum_<cutlass::bfloat16_t>(FWD_PARAMS);
+  if (!kIsMeta) {
+    if (a.scalar_type() == at::ScalarType::Half) {
+      TORCH_CHECK(b.scalar_type() == at::ScalarType::Half);
+      TORCH_CHECK(out_mm.scalar_type() == at::ScalarType::Half);
+      TORCH_CHECK(out_sum.scalar_type() == at::ScalarType::Half);
+      gemm_fused_operand_sum_<cutlass::half_t>(FWD_PARAMS);
+    } else {
+      TORCH_CHECK(
+          a.scalar_type() == at::ScalarType::BFloat16,
+          "Only supports bf16/f16");
+      TORCH_CHECK(b.scalar_type() == at::ScalarType::BFloat16);
+      TORCH_CHECK(out_mm.scalar_type() == at::ScalarType::BFloat16);
+      TORCH_CHECK(out_sum.scalar_type() == at::ScalarType::BFloat16);
+      gemm_fused_operand_sum_<cutlass::bfloat16_t>(FWD_PARAMS);
+    }
   }
   return std::make_tuple(out_mm, out_sum);
 }
@@ -251,7 +255,13 @@ std::tuple<at::Tensor, at::Tensor> gemm_fused_operand_sum_autocast(
 TORCH_LIBRARY_IMPL(xformers, CUDA, m) {
   m.impl(
       TORCH_SELECTIVE_NAME("xformers::gemm_fused_operand_sum"),
-      TORCH_FN(gemm_fused_operand_sum));
+      TORCH_FN(gemm_fused_operand_sum<false>));
+}
+
+TORCH_LIBRARY_IMPL(xformers, Meta, m) {
+  m.impl(
+      TORCH_SELECTIVE_NAME("xformers::gemm_fused_operand_sum"),
+      TORCH_FN(gemm_fused_operand_sum<true>));
 }
 
 TORCH_LIBRARY_IMPL(xformers, Autocast, m) {

{xformers-0.0.24.dev720 → xformers-0.0.24.dev725}/xformers/csrc/swiglu/cuda/silu_bw_fused.cu

@@ -9,6 +9,7 @@
 #include <ATen/Dispatch.h>
 #include <ATen/ScalarOps.h>
 #include <ATen/Tensor.h>
+#include <ATen/autocast_mode.h>
 #include <ATen/cuda/CUDAContext.h>
 #include <ATen/native/ReduceOps.h>
 #include <ATen/native/Resize.h>
@@ -32,6 +33,7 @@ def silu_bw_fused(x1, x2, dx4):
     return dx1, dx2, x4
 */
 
+template <bool kIsMeta = false>
 std::tuple<at::Tensor, at::Tensor> silu_bw_fused(
     const at::Tensor& x1,
     const at::Tensor& x2,
@@ -39,15 +41,26 @@ std::tuple<at::Tensor, at::Tensor> silu_bw_fused(
   // TODO: Check all params. This would take a lot of lines of code...
   TORCH_CHECK(x2.dim() == 2);
   TORCH_CHECK(dx4.dim() == 2);
-  TORCH_CHECK(x2.size(0) == dx4.size(0));
-  TORCH_CHECK(x2.size(1) == dx4.size(1));
+  TORCH_CHECK(x2.sym_size(0) == dx4.sym_size(0));
+  TORCH_CHECK(x2.sym_size(1) == dx4.sym_size(1));
 
-  int64_t B = x2.size(0);
-  int64_t H = x2.size(1);
-  at::Tensor dx1dx2 = at::empty({B, 2, H}, x2.options());
+  at::SymInt B = x2.sym_size(0);
+  at::SymInt H = x2.sym_size(1);
+  at::Tensor dx1dx2 = at::empty_symint({B, 2, H}, x2.options());
+  at::Tensor x4 = at::empty_symint({B, H}, x2.options());
+
+  // Check if the function is in meta mode
+  if (kIsMeta) {
+    // Meta mode logic: return tensors with appropriate shapes
+    // Infer the shapes based on input tensors and return empty tensors with
+    // those shapes
+    return std::make_tuple(dx1dx2, x4);
+  }
+
+  // Regular mode logic: perform actual computations
   at::Tensor dx1 = dx1dx2.select(1, 0);
   at::Tensor dx2 = dx1dx2.select(1, 1);
-  at::Tensor x4 = at::empty({B, H}, x2.options());
+
   auto iter = at::TensorIteratorConfig()
                   .add_output(dx1)
                   .add_output(dx2)
@@ -83,9 +96,34 @@ std::tuple<at::Tensor, at::Tensor> silu_bw_fused(
       }));
   return std::make_tuple(dx1dx2, x4);
 }
+
+std::tuple<at::Tensor, at::Tensor> silu_bw_fused_autocast(
+    const at::Tensor& x1,
+    const at::Tensor& x2,
+    const at::Tensor& dx4) {
+  c10::impl::ExcludeDispatchKeyGuard no_autocast(c10::DispatchKey::Autocast);
+  auto exec_type = at::autocast::get_autocast_gpu_dtype();
+  return silu_bw_fused(
+      at::autocast::cached_cast(exec_type, x1),
+      at::autocast::cached_cast(exec_type, x2),
+      at::autocast::cached_cast(exec_type, dx4));
+}
 } // namespace
 
 TORCH_LIBRARY_IMPL(xformers, CUDA, m) {
   m.impl(
-      TORCH_SELECTIVE_NAME("xformers::silu_bw_fused"), TORCH_FN(silu_bw_fused));
+      TORCH_SELECTIVE_NAME("xformers::silu_bw_fused"),
+      TORCH_FN(silu_bw_fused<false>));
+}
+
+TORCH_LIBRARY_IMPL(xformers, Meta, m) {
+  m.impl(
+      TORCH_SELECTIVE_NAME("xformers::silu_bw_fused"),
+      TORCH_FN(silu_bw_fused<true>));
+}
+
+TORCH_LIBRARY_IMPL(xformers, Autocast, m) {
+  m.impl(
+      TORCH_SELECTIVE_NAME("xformers::silu_bw_fused"),
+      TORCH_FN(silu_bw_fused_autocast));
 }

{xformers-0.0.24.dev720 → xformers-0.0.24.dev725}/xformers/csrc/swiglu/swiglu_packedw.cpp

@@ -236,3 +236,7 @@ TORCH_LIBRARY_IMPL(xformers, Autocast, m) {
 TORCH_LIBRARY_IMPL(xformers, CUDA, m) {
   m.impl("swiglu_packedw", swiglu_packedw_cuda);
 }
+
+TORCH_LIBRARY_IMPL(xformers, Meta, m) {
+  m.impl("swiglu_packedw", swiglu_packedw_cuda);
+}

{xformers-0.0.24.dev720 → xformers-0.0.24.dev725}/xformers/ops/_triton/rmsnorm_kernels.py

@@ -19,7 +19,7 @@ def _rms_norm_kernel(
     w_ptr,
     eps,
     stride,
-    N_COLS,
+    N_COLS: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
     INCLUDE_WEIGHT: tl.constexpr,
 ):
@@ -56,7 +56,7 @@ def _rms_norm_add_kernel(
     w_ptr,
     eps,
     stride,
-    N_COLS,
+    N_COLS: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
     INCLUDE_WEIGHT: tl.constexpr,
 ):
@@ -105,7 +105,7 @@ def _rms_norm_forward(x, attn_norm_weights, eps):
     MAX_FUSED_SIZE = 65536 // x.element_size()
     BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(N))
     BLOCK_SIZE = max(BLOCK_SIZE, 128)
-    BLOCK_SIZE = min(BLOCK_SIZE, 4096)
+    BLOCK_SIZE = min(BLOCK_SIZE, 8192)
     # heuristics for number of warps
     num_warps = min(max(BLOCK_SIZE // 256, 1), 8)
     _rms_norm_kernel[(M,)](
@@ -140,7 +140,7 @@ def _rms_norm_add_forward(x, y, attn_norm_weights, eps):
     MAX_FUSED_SIZE = 65536 // x.element_size()
     BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(N))
     BLOCK_SIZE = max(BLOCK_SIZE, 128)
-    BLOCK_SIZE = min(BLOCK_SIZE, 4096)
+    BLOCK_SIZE = min(BLOCK_SIZE, 8192)
     # heuristics for number of warps
     num_warps = min(max(BLOCK_SIZE // 256, 1), 8)
     _rms_norm_add_kernel[(M,)](