xformers 0.0.24.dev728__tar.gz → 0.0.24.dev734__tar.gz

{xformers-0.0.24.dev728/xformers.egg-info → xformers-0.0.24.dev734}/PKG-INFO

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

xformers-0.0.24.dev734/version.txt

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

{xformers-0.0.24.dev728 → xformers-0.0.24.dev734}/xformers/__init__.py

@@ -47,6 +47,10 @@ def _is_triton_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:
         from xformers.triton.softmax import softmax as triton_softmax  # noqa
 

{xformers-0.0.24.dev728 → xformers-0.0.24.dev734}/xformers/components/attention/blocksparse.py

@@ -23,18 +23,6 @@ if _is_blocksparse_available:
     from triton.ops.blocksparse import matmul as blocksparse_matmul  # type: ignore
     from triton.ops.blocksparse import softmax as blocksparse_softmax  # type: ignore
 
-    from xformers.triton.utils import gpu_capabilities_older_than_80
-
-    # Blocksparse requires Tensor cores, but we also disable it on V100 because of Triton issues
-    if gpu_capabilities_older_than_80():
-        logger.warning(
-            "Blocksparse is not available: the current GPU does not expose Tensor cores"
-        )
-        _is_blocksparse_available = False
-
-
-if _is_blocksparse_available:
-
     @dataclass
     class BlockSparseAttentionConfig(AttentionConfig):
         layout: torch.Tensor  # The dimensions of the random features

{xformers-0.0.24.dev728 → xformers-0.0.24.dev734}/xformers/components/attention/core.py

@@ -18,16 +18,10 @@ from xformers.components.attention.attention_mask import AttentionMask
 if _has_cpp_library:
     from ._sputnik_sparse import SparseCS
 
-if _is_triton_available():
-    from xformers.triton.softmax import softmax as triton_softmax
-    from xformers.triton.utils import gpu_capabilities_older_than_80
-
-_is_blocksparse_available = (
-    _is_triton_available() and not gpu_capabilities_older_than_80()
-)
-
+_is_blocksparse_available = _is_triton_available()
 if _is_blocksparse_available:
     from xformers.components.attention.blocksparse import BlockSparseAttention
+    from xformers.triton.softmax import softmax as triton_softmax
 
 
 logger = logging.getLogger("xformers")

{xformers-0.0.24.dev728 → xformers-0.0.24.dev734}/xformers/ops/_triton/sequence_parallel_fused_kernels.py

@@ -130,7 +130,18 @@ def determine_tile(
     world_size,
     direction,
     stride_am,
-    stride_cm,
+    stride_bk1,
+    stride_bk2,
+    stride_bk3,
+    stride_bn1,
+    stride_bn2,
+    stride_bn3,
+    stride_cm1,
+    stride_cm2,
+    stride_cm3,
+    stride_cn1,
+    stride_cn2,
+    stride_cn3,
     BLOCK_M,
     BLOCK_N,
     GROUP_M,
@@ -173,6 +184,26 @@ def determine_tile(
         C1_my_shard,
         tl.where(pid_n < grid_n1 + grid_n2, C2_my_shard, C3_my_shard),
     )
+    stride_bk = tl.where(
+        pid_n < grid_n1,
+        stride_bk1,
+        tl.where(pid_n < grid_n1 + grid_n2, stride_bk2, stride_bk3),
+    )
+    stride_bn = tl.where(
+        pid_n < grid_n1,
+        stride_bn1,
+        tl.where(pid_n < grid_n1 + grid_n2, stride_bn2, stride_bn3),
+    )
+    stride_cm = tl.where(
+        pid_n < grid_n1,
+        stride_cm1,
+        tl.where(pid_n < grid_n1 + grid_n2, stride_cm2, stride_cm3),
+    )
+    stride_cn = tl.where(
+        pid_n < grid_n1,
+        stride_cn1,
+        tl.where(pid_n < grid_n1 + grid_n2, stride_cn2, stride_cn3),
+    )
     N = tl.where(pid_n < grid_n1, N1, tl.where(pid_n < grid_n1 + grid_n2, N2, N3))
     pid_n = tl.where(
         pid_n < grid_n1,
@@ -187,7 +218,21 @@ def determine_tile(
         other_rank == my_rank, C_my_shard, C + other_rank * M_per_rank * stride_cm
     )
 
-    return A, B, C, M_per_rank, N, pid_m, pid_n, other_rank, blocks_per_rank
+    return (
+        A,
+        B,
+        C,
+        M_per_rank,
+        N,
+        stride_bk,
+        stride_bn,
+        stride_cm,
+        stride_cn,
+        pid_m,
+        pid_n,
+        other_rank,
+        blocks_per_rank,
+    )
 
 
 @triton.jit
@@ -338,7 +383,21 @@ def our_estimate_matmul_time(B1, C1, N1, N2, N3, **kwargs):
     }
 )
 @triton.jit(
-    do_not_specialize=[11, 12, 13, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34],
+    do_not_specialize=[
+        "wait_counters",
+        "blocks_done_counters",
+        "write_counters",
+        "do_wait",
+        "do_write",
+        "direction",
+        "stripe",
+        "seq_num",
+        "num_stripes",
+        "_wait",
+        "my_rank",
+        "world_size",
+        "timeout_ns",
+    ],
     debug=True,  # To avoid stripping device asserts
 )
 def _xformers_seqpar_matmul_kernel(
@@ -363,10 +422,18 @@ def _xformers_seqpar_matmul_kernel(
     K,  # DO NOT CHANGE NAME: MUST MATCH PERF MODEL
     stride_am,
     stride_ak,
-    stride_bk,
-    stride_bn,
-    stride_cm,
-    stride_cn,
+    stride_bk1,
+    stride_bk2,
+    stride_bk3,
+    stride_bn1,
+    stride_bn2,
+    stride_bn3,
+    stride_cm1,
+    stride_cm2,
+    stride_cm3,
+    stride_cn1,
+    stride_cn2,
+    stride_cn3,
     do_wait,
     do_write,
     direction,
@@ -385,7 +452,21 @@ def _xformers_seqpar_matmul_kernel(
     EVEN_K: tl.constexpr,
     ACC_TYPE: tl.constexpr,
 ):
-    A, B, C, M, N, pid_m, pid_n, other_rank, num_blocks_2d = determine_tile(
+    (
+        A,
+        B,
+        C,
+        M,
+        N,
+        stride_bk,
+        stride_bn,
+        stride_cm,
+        stride_cn,
+        pid_m,
+        pid_n,
+        other_rank,
+        num_blocks_2d,
+    ) = determine_tile(
         A,
         B1,
         B2,
@@ -405,7 +486,18 @@ def _xformers_seqpar_matmul_kernel(
         world_size,
         direction,
         stride_am,
-        stride_cm,
+        stride_bk1,
+        stride_bk2,
+        stride_bk3,
+        stride_bn1,
+        stride_bn2,
+        stride_bn3,
+        stride_cm1,
+        stride_cm2,
+        stride_cm3,
+        stride_cn1,
+        stride_cn2,
+        stride_cn3,
         BLOCK_M,
         BLOCK_N,
         GROUP_M,
@@ -504,13 +596,11 @@ def _launch_triton_matmul(
     assert all(c.shape[0] == M for c in cs)
     assert all(c.shape[1] == N for c, N in zip(cs, Ns))
     stride_am, stride_ak = cast(Tuple[int, int], a.stride())
-    stride_bk, stride_bn = cast(Tuple[int, int], bs[0].stride())
-    stride_cm, stride_cn = cast(Tuple[int, int], cs[0].stride())
-    assert all(b.stride() == (stride_bk, stride_bn) for b in bs)
-    assert all(c.stride() == (stride_cm, stride_cn) for c in cs)
+    strides_bk, strides_bn = zip(*(cast(Tuple[int, int], b.stride()) for b in bs))
+    strides_cm, strides_cn = zip(*(cast(Tuple[int, int], c.stride()) for c in cs))
     assert stride_am == 1 or stride_ak == 1
-    assert stride_bk == 1 or stride_bn == 1
-    assert stride_cm == 1 or stride_cn == 1
+    assert all(s == 1 for s in strides_bk) or all(s == 1 for s in strides_bn)
+    assert all(s == 1 for s in strides_cm) or all(s == 1 for s in strides_cn)
 
     if a_my_shard is not None:
         assert a_my_shard.ndim == 2
@@ -615,10 +705,18 @@ def _launch_triton_matmul(
     kwargs = dict(
         stride_am=stride_am,
         stride_ak=stride_ak,
-        stride_bk=stride_bk,
-        stride_bn=stride_bn,
-        stride_cm=stride_cm,
-        stride_cn=stride_cn,
+        stride_bk1=strides_bk[0],
+        stride_bk2=strides_bk[min(1, len(strides_bk) - 1)],
+        stride_bk3=strides_bk[min(2, len(strides_bk) - 1)],
+        stride_bn1=strides_bn[0],
+        stride_bn2=strides_bn[min(1, len(strides_bn) - 1)],
+        stride_bn3=strides_bn[min(2, len(strides_bn) - 1)],
+        stride_cm1=strides_cm[0],
+        stride_cm2=strides_cm[min(1, len(strides_cm) - 1)],
+        stride_cm3=strides_cm[min(2, len(strides_cm) - 1)],
+        stride_cn1=strides_cn[0],
+        stride_cn2=strides_cn[min(1, len(strides_cn) - 1)],
+        stride_cn3=strides_cn[min(2, len(strides_cn) - 1)],
         do_wait=do_wait,
         do_write=do_write,
         direction=direction,

{xformers-0.0.24.dev728 → xformers-0.0.24.dev734}/xformers/ops/common.py

@@ -4,7 +4,6 @@
 # LICENSE file in the root directory of this source tree.
 
 import inspect
-import os
 from dataclasses import dataclass
 from functools import wraps
 from typing import Any, Callable, Dict, List, Type, TypeVar, Union
@@ -13,6 +12,8 @@ import torch
 from torch.torch_version import TorchVersion
 from typing_extensions import Annotated, get_args, get_origin
 
+from .. import _is_triton_available
+
 
 def get_operator(library: str, name: str):
     def no_such_operator(*args, **kwargs):
@@ -167,20 +168,8 @@ def turn_into_pytorch_op(fn: ClsT, dispatch_key: str) -> ClsT:
     return caller  # type: ignore
 
 
-def _has_a_version_of_triton():
-    if os.environ.get("XFORMERS_FORCE_DISABLE_TRITON", "0") == "1":
-        return False
-    if not torch.cuda.is_available():
-        return False
-    try:
-        import triton  # noqa: F401
-    except ImportError:
-        return False
-    return True
-
-
 def _has_triton2():
-    if not _has_a_version_of_triton():
+    if not _is_triton_available():
         return False
     import triton
 
@@ -189,7 +178,7 @@ def _has_triton2():
 
 
 def _has_triton21():
-    if not _has_a_version_of_triton():
+    if not _is_triton_available():
         return False
     import triton
 

{xformers-0.0.24.dev728 → xformers-0.0.24.dev734}/xformers/ops/differentiable_collectives.py

@@ -15,7 +15,7 @@ def all_reduce(
 ) -> None:
     assert x.is_contiguous()
 
-    mp_size = torch.distributed.get_world_size(process_group)
+    mp_size = process_group.size()
     if mp_size == 1:
         return
 
@@ -28,7 +28,7 @@ def gather_along_first_dim_async(
     input_: torch.Tensor, *, process_group: torch.distributed.ProcessGroup
 ) -> Tuple[torch.Tensor, Optional[torch.distributed.Work]]:
     assert input_.is_contiguous()
-    mp_size = torch.distributed.get_world_size(process_group)
+    mp_size = process_group.size()
     if mp_size == 1:
         return input_, None
 
@@ -47,7 +47,7 @@ def reduce_scatter_along_first_dim_async(
     input_: torch.Tensor, *, process_group: torch.distributed.ProcessGroup
 ) -> Tuple[torch.Tensor, Optional[torch.distributed.Work]]:
     assert input_.is_contiguous()
-    mp_size = torch.distributed.get_world_size(process_group)
+    mp_size = process_group.size()
     if mp_size == 1:
         return input_, None
 

{xformers-0.0.24.dev728 → xformers-0.0.24.dev734}/xformers/ops/modpar_layers.py

@@ -25,8 +25,8 @@ def _init_2d_weight(
     # The reason we initialize the full unpartitioned/gathered weight is so that
     # different ranks get different initial values and thus "break the symmetry"
     # and in order to achieve the same init for any value of model parallelism.
-    rank = torch.distributed.get_rank(process_group)
-    world_size = torch.distributed.get_world_size(process_group)
+    rank = process_group.rank()
+    world_size = process_group.size()
 
     nrows, ncols = weight.shape
     if partition_dim == 0:
@@ -77,7 +77,7 @@ class ColumnParallelLinear(torch.nn.Module):
         self.sequence_parallel = sequence_parallel
         self.fuse_sequence_parallel = fuse_sequence_parallel
         self.process_group = process_group
-        mp_size = torch.distributed.get_world_size(process_group)
+        mp_size = process_group.size()
         assert all(dim % mp_size == 0 for dim in out_features)
         self.my_out_features = [dim // mp_size for dim in out_features]
 
@@ -136,7 +136,7 @@ class RowParallelLinear(torch.nn.Module):
         self.sequence_parallel = sequence_parallel
         self.fuse_sequence_parallel = fuse_sequence_parallel
         self.process_group = process_group
-        mp_size = torch.distributed.get_world_size(process_group)
+        mp_size = process_group.size()
         assert in_features % mp_size == 0
         self.my_in_features = in_features // mp_size
 

{xformers-0.0.24.dev728 → xformers-0.0.24.dev734}/xformers/ops/seqpar.py

@@ -54,7 +54,7 @@ def sequence_parallel_leading_matmul_bwd(
     fuse: bool,
     process_group: torch.distributed.ProcessGroup,
 ) -> Tuple[torch.Tensor, List[torch.Tensor]]:
-    mp_size = torch.distributed.get_world_size(process_group)
+    mp_size = process_group.size()
 
     if fuse:
         grad_scattered_input = torch.empty_like(scattered_input)
@@ -206,7 +206,7 @@ def sequence_parallel_trailing_matmul_bwd(
     fuse: bool,
     process_group: torch.distributed.ProcessGroup,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
-    mp_size = torch.distributed.get_world_size(process_group)
+    mp_size = process_group.size()
 
     if fuse:
         grad_gathered_input = torch.empty_like(gathered_input)