[fmha] LowerTriangularFromBottomRightLocalAttentionMask

ghstack-source-id: a7ff2da
Pull Request resolved: https://github.com/fairinternal/xformers/pull/933

__original_commit__ = fairinternal/xformers@9801b557e356e5c032f5af064e00f22dafe164bf

Author: xFormers Bot

docs/source/_static/causal_bottom_right_local.png

@@ -43,6 +43,7 @@ def create_attn_bias(
     if bias_type is None or isinstance(None, bias_type):
         return None
     r = random.Random("-".join(map(str, [batch_size, q_len, kv_len, dtype, fmt])))
+    window_size = {0: 3, 1: 128, 2: 300}[r.randint(0, 2)]
     if bias_type is torch.Tensor:
         if fmt == "BMK":
             batch_size *= num_heads
@@ -81,6 +82,8 @@ def create_attn_bias(
         return bias_type()
     if bias_type is fmha.attn_bias.LowerTriangularFromBottomRightMask:
         return bias_type()
+    if bias_type is fmha.attn_bias.LowerTriangularFromBottomRightLocalAttentionMask:
+        return bias_type(window_size)
     if bias_type is fmha.attn_bias.LowerTriangularMaskWithTensorBias:
         attn_bias = _create_aligned_bias(
             batch_size,
@@ -107,13 +110,6 @@ def create_attn_bias(
     ]:
         # These bias types are not supported in BMK format
         assert fmt in ["BMGHK", "BMHK"]
-
-        window_size = None
-        if bias_type in {
-            fmha.attn_bias.BlockDiagonalCausalLocalAttentionMask,
-            fmha.attn_bias.BlockDiagonalCausalLocalAttentionFromBottomRightMask,
-        }:
-            window_size = {0: 3, 1: 128, 2: 300}[r.randint(0, 2)]
         max_q_minus_k = None
         if bias_type in {
             fmha.attn_bias.BlockDiagonalCausalFromBottomRightMask,

xformers/attn_bias_utils.py

@@ -112,6 +112,9 @@ def materialize(
         return _materialize_causal_mask(shape, dtype=dtype, device=device)
 
     def add_bias(self, bias: torch.Tensor) -> "LowerTriangularMaskWithTensorBias":
+        """
+        Creates a new causal mask with an arbitrary ``torch.Tensor`` bias
+        """
         return LowerTriangularMaskWithTensorBias(bias)
 
 
@@ -145,6 +148,60 @@ def materialize(
             shape, dtype=dtype, device=device, from_bottomright=True
         )
 
+    def make_local_attention(
+        self, window_size: int
+    ) -> "LowerTriangularFromBottomRightLocalAttentionMask":
+        """
+        Create a new bias which combines local + causal attention.
+
+        See :attr:`LowerTriangularFromBottomRightLocalAttentionMask`
+        """
+        return LowerTriangularFromBottomRightLocalAttentionMask(window_size)
+
+
+@dataclass
+class LowerTriangularFromBottomRightLocalAttentionMask(
+    LowerTriangularFromBottomRightMask
+):
+    """
+    A mask that combines both :attr:`LowerTriangularFromBottomRightMask` and
+    local attention.
+
+    A query whose distance from the final query is X cannot attend to a key
+    whose distance to the final key is either of:
+
+    * less than X (i.e. "causal attention", same as :attr:`LowerTriangularFromBottomRightMask`)
+    * greater than X + window_size (i.e. "local attention")
+
+
+    .. figure:: /_static/causal_bottom_right_local.png
+
+        The mask from :attr:`LowerTriangularFromBottomRightLocalAttentionMask`.
+        The green area is calculated, and the grey area is masked out.
+    """
+
+    _window_size: int
+
+    def __post_init__(self) -> None:
+        if self._window_size <= 0:
+            raise ValueError(
+                f"Expected `window_size > 0`, but window_size={self._window_size}"
+            )
+
+    def materialize(
+        self,
+        shape: Tuple[int, ...],
+        dtype: torch.dtype = torch.float32,
+        device: Union[str, torch.device] = "cpu",
+    ) -> torch.Tensor:
+        return _materialize_causal_mask(
+            shape,
+            dtype=dtype,
+            device=device,
+            window_size=self._window_size,
+            from_bottomright=True,
+        )
+
 
 class LowerTriangularMaskWithTensorBias(LowerTriangularMask):
     """A lower-triangular (aka causal) mask with an additive bias"""

xformers/ops/fmha/attn_bias.py

@@ -400,7 +400,7 @@ def attn_operator_flop(
 
 class AttentionBwOpBase(AttentionOpBase):
     ERROR_ATOL: Mapping[torch.dtype, float] = {
-        torch.float: 5e-4,
+        torch.float: 6e-4,
         torch.half: 9e-2,
         torch.bfloat16: 0.7,
     }

xformers/ops/fmha/common.py

@@ -20,6 +20,7 @@
     BlockDiagonalCausalMask,
     BlockDiagonalCausalWithOffsetPaddedKeysMask,
     BlockDiagonalMask,
+    LowerTriangularFromBottomRightLocalAttentionMask,
     LowerTriangularFromBottomRightMask,
     LowerTriangularMask,
     LowerTriangularMaskWithTensorBias,
@@ -144,6 +145,7 @@ def _custom_mask_type(bias: Optional[Union[torch.Tensor, AttentionBias]]) -> int
         bias,
         (
             LowerTriangularFromBottomRightMask,
+            LowerTriangularFromBottomRightLocalAttentionMask,
             attn_bias.BlockDiagonalCausalFromBottomRightMask,
             BlockDiagonalCausalWithOffsetPaddedKeysMask,
             BlockDiagonalCausalLocalAttentionFromBottomRightMask,
@@ -169,6 +171,7 @@ class FwOp(AttentionFwOpBase):
         torch.Tensor,
         LowerTriangularMask,
         LowerTriangularFromBottomRightMask,
+        LowerTriangularFromBottomRightLocalAttentionMask,
         LowerTriangularMaskWithTensorBias,
         BlockDiagonalMask,
         BlockDiagonalCausalMask,
@@ -281,6 +284,7 @@ def apply_bmhk(
                 (
                     BlockDiagonalCausalLocalAttentionMask,
                     BlockDiagonalCausalLocalAttentionFromBottomRightMask,
+                    LowerTriangularFromBottomRightLocalAttentionMask,
                 ),
             )
             else None,
@@ -348,6 +352,8 @@ class BwOp(AttentionBwOpBase):
         torch.Tensor,
         LowerTriangularMask,
         LowerTriangularFromBottomRightMask,
+        # TODO: Still some bugs in the BW pass for that one
+        # LowerTriangularFromBottomRightLocalAttentionMask,
         # TODO: Fix handling of gradient through the fMHA autograd function
         # LowerTriangularMaskWithTensorBias,
         BlockDiagonalMask,
@@ -455,6 +461,7 @@ def apply(cls, ctx: Context, inp: Inputs, grad: torch.Tensor) -> Gradients:
                 (
                     BlockDiagonalCausalLocalAttentionMask,
                     BlockDiagonalCausalLocalAttentionFromBottomRightMask,
+                    LowerTriangularFromBottomRightLocalAttentionMask,
                 ),
             )
             else None,

xformers/ops/fmha/cutlass.py

@@ -18,6 +18,7 @@
     BlockDiagonalCausalLocalAttentionMask,
     BlockDiagonalCausalMask,
     BlockDiagonalMask,
+    LowerTriangularFromBottomRightLocalAttentionMask,
     LowerTriangularFromBottomRightMask,
     LowerTriangularMask,
 )
@@ -285,6 +286,7 @@ def _is_causal(attn_bias: Optional[Union[torch.Tensor, AttentionBias]]) -> bool:
         (
             LowerTriangularMask,
             LowerTriangularFromBottomRightMask,
+            LowerTriangularFromBottomRightLocalAttentionMask,
             BlockDiagonalCausalMask,
             BlockDiagonalCausalLocalAttentionMask,
             BlockDiagonalCausalFromBottomRightMask,
@@ -301,6 +303,8 @@ def _window_size(attn_bias: Optional[Union[torch.Tensor, AttentionBias]]) -> int
         return attn_bias._window_size or 0
     if isinstance(attn_bias, BlockDiagonalCausalLocalAttentionFromBottomRightMask):
         return attn_bias._window_size
+    if isinstance(attn_bias, LowerTriangularFromBottomRightLocalAttentionMask):
+        return attn_bias._window_size
     return 0
 
 
@@ -360,6 +364,7 @@ class FwOp(AttentionFwOpBase):
         type(None),
         LowerTriangularMask,
         LowerTriangularFromBottomRightMask,
+        LowerTriangularFromBottomRightLocalAttentionMask,
         BlockDiagonalMask,
         BlockDiagonalCausalMask,
         BlockDiagonalCausalLocalAttentionMask,