xformers 0.0.25.dev745__tar.gz → 0.0.25.dev746__tar.gz

{xformers-0.0.25.dev745/xformers.egg-info → xformers-0.0.25.dev746}/PKG-INFO

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

xformers-0.0.25.dev746/version.txt

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

{xformers-0.0.25.dev745 → xformers-0.0.25.dev746}/xformers/attn_bias_utils.py

@@ -150,7 +150,10 @@ def create_attn_bias(
         if bias_type is fmha.attn_bias.BlockDiagonalCausalFromBottomRightMask:
             block_diag = block_diag.make_causal_from_bottomright()
         return block_diag
-    if bias_type == fmha.attn_bias.BlockDiagonalCausalWithOffsetPaddedKeysMask:
+    if bias_type in [
+        fmha.attn_bias.BlockDiagonalCausalWithOffsetPaddedKeysMask,
+        fmha.attn_bias.PagedBlockDiagonalCausalWithOffsetPaddedKeysMask,
+    ]:
         assert fmt in ["BMHK", "BMGHK"]
         q, k = _rand_seqlens_padded_k(r, batch_size, q_len, kv_len)
         g_block_diag = (
@@ -160,6 +163,15 @@ def create_attn_bias(
                 kv_seqlen=k,
             )
         )
+        if bias_type == fmha.attn_bias.PagedBlockDiagonalCausalWithOffsetPaddedKeysMask:
+            page_size = r.choice([64, 128, 256])
+            pages_per_row = (kv_len + page_size - 1) // page_size
+            block_tables = torch.randperm(
+                batch_size * pages_per_row, device=device
+            ).reshape(batch_size, pages_per_row)
+            return g_block_diag.make_paged(
+                block_tables=block_tables, page_size=page_size
+            )
         return g_block_diag
     if bias_type == fmha.attn_bias.LocalAttentionFromBottomRightMask:
         return bias_type(

{xformers-0.0.25.dev745 → xformers-0.0.25.dev746}/xformers/benchmarks/benchmark_attn_decoding.py

@@ -67,7 +67,10 @@ class AttentionDecodingFlashDecoding:
         self, B: int, Mq: int, Mkv: int, Hq: int, Hkv: int, K: int, bw: bool
     ) -> None:
         dtype = torch.float16
-        self.sub_label = f"B={B} Mq={Mq} Mkv={Mkv} Hq={Hq} Hkv={Hkv} K={K}"
+        self.sub_label = (
+            f"B={B} Mq={Mq} Mkv={Mkv} Hq={Hq} Hkv={Hkv} K={K} TotalBytes="
+            f"{((B * Mkv * Hkv * K * 2) + (B * Mq * Hq * K) + (B * Mq * Hq * K)) * 2}"
+        )
         self.label = "attn_decoding"
         self.shapes = (B, Mq, Mkv, Hq, Hkv, K)
 

{xformers-0.0.25.dev745 → xformers-0.0.25.dev746}/xformers/ops/fmha/__init__.py

@@ -414,6 +414,68 @@ def _memory_efficient_attention_backward(
     return grads
 
 
+def merge_attentions(
+    attn_split: torch.Tensor, lse_split: torch.Tensor, write_lse: bool = True
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    """
+    Combine attention output computed on different parts of K/V for the same
+    query to get attention on the whole K/V. See https://arxiv.org/abs/2402.05099
+    The result is equal to
+        Out_full = (Out1 * exp(LSE1) + Out2 * exp(LSE2) + ...) / (exp(LSE1) + exp(LSE2) + ...)
+        LSE_full = log(exp(LSE1) + exp(LSE2) + ...)
+    Attention inputs are in BH(G)MK format, stacked along dim 0. Attention output also is in BH(G)MK.
+    Args:
+        attn_split: [split_k, B, H, G, M, Kq] or [split_k, B, H, M, Kq]
+        lse_split: [split_k, B, H, G, M] or [split_k, B, H, M]
+    Res:
+        attn_out: [B, H, G, M, K] or [B, H, M, K]
+        lse_out: [B, H, G, M] or [B, H, M]
+    """
+
+    assert (
+        attn_split.ndim == lse_split.ndim + 1
+    ), f"{attn_split.shape=} {lse_split.shape=}"
+
+    is_bmhk = attn_split.ndim == 5
+    if is_bmhk:
+        attn_split = attn_split.unsqueeze(3)
+        lse_split = lse_split.unsqueeze(3)
+
+    split_k, B, H, G, M_ceil, Kq = attn_split.shape
+    split_k1, B1, H1, G1, M = lse_split.shape
+    assert (
+        B == B1 and G == G1 and H == H1 and split_k == split_k1 and M_ceil >= M
+    ), f"{attn_split.shape=} {lse_split.shape=}"
+
+    attn_split = attn_split.permute(1, 2, 3, 0, 4, 5).view(
+        B, H * G, split_k, M_ceil, Kq
+    )
+    lse_split = lse_split.permute(1, 2, 3, 0, 4).view(B, H * G, split_k, M)
+
+    attn_out = torch.empty(
+        B, H, G, M, Kq, device=attn_split.device, dtype=attn_split.dtype
+    )
+    if write_lse:
+        lse_out = torch.empty(
+            B * H * G, M, device=attn_split.device, dtype=torch.float32
+        )
+    else:
+        lse_out = None
+
+    triton_splitk.merge_attentions(
+        attn_out.permute(0, 3, 2, 1, 4), lse_out, attn_split, lse_split
+    )
+    if lse_out is not None:
+        lse_out = lse_out.view(B, H, G, M)
+
+    if is_bmhk:
+        attn_out = attn_out[:, :, 0]
+        if lse_out is not None:
+            lse_out = lse_out[:, :, 0]
+
+    return attn_out, lse_out
+
+
 ALL_FW_OPS: Sequence[Type[AttentionFwOpBase]] = [
     cutlass.FwOp,
     flash.FwOp,

{xformers-0.0.25.dev745 → xformers-0.0.25.dev746}/xformers/ops/fmha/attn_bias.py

@@ -327,10 +327,14 @@ class _SeqLenInfo:
         yield from zip(self.seqstart_py, self.seqstart_py[1:])
 
     @classmethod
-    def from_seqlens(cls, seqlens: Iterable[int]) -> "_SeqLenInfo":
+    def _get_seqstart(
+        cls, seqlens: Iterable[int]
+    ) -> Tuple[int, int, List[int], torch.Tensor]:
         """
-        Input tensors are assumed to be in shape [B, M, *]
+        Given sequence lengths, returns the min/max value and the sequence start
+        positions (offsets), with first element being 0 (returned in list and Tensor).
         """
+
         assert not isinstance(seqlens, torch.Tensor)
         seqstart_py = [0]
         max_seqlen = -1
@@ -340,6 +344,16 @@ class _SeqLenInfo:
             max_seqlen = max(max_seqlen, seqlen)
             seqstart_py.append(seqstart_py[len(seqstart_py) - 1] + seqlen)
         seqstart = torch.tensor(seqstart_py, dtype=torch.int32)
+
+        return (min_seqlen, max_seqlen, seqstart_py, seqstart)
+
+    @classmethod
+    def from_seqlens(cls, seqlens: Iterable[int]) -> "_SeqLenInfo":
+        """
+        Input tensors are assumed to be in shape [B, M, *]
+        """
+        min_seqlen, max_seqlen, seqstart_py, seqstart = cls._get_seqstart(seqlens)
+
         return cls(
             max_seqlen=max_seqlen,
             min_seqlen=min_seqlen,
@@ -440,7 +454,9 @@ class _PaddedSeqLenInfo(_SeqLenInfo):
         seqstart = padding * torch.arange(batch_size)
         """
         assert not isinstance(seqlens, torch.Tensor)
-        assert all(seqlen <= padding for seqlen in seqlens)
+        assert all(
+            seqlen <= padding for seqlen in seqlens
+        ), f"Seqlens {seqlens} Padding {padding}"
         seqstart_py = list(range(0, len(seqlens) * padding + 1, padding))
         return cls(
             seqlen=torch.tensor(seqlens, dtype=torch.int32),
@@ -839,6 +855,106 @@ class BlockDiagonalCausalWithOffsetPaddedKeysMask(AttentionBias):
         k_seqinfo = _PaddedSeqLenInfo.from_seqlens_padded(kv_seqlen, kv_padding)
         return cls(q_seqinfo=q_seqinfo, k_seqinfo=k_seqinfo)
 
+    def make_paged(self, block_tables: torch.Tensor, page_size: int):
+        paged_bias = PagedBlockDiagonalCausalWithOffsetPaddedKeysMask(
+            q_seqinfo=self.q_seqinfo,
+            k_seqinfo=self.k_seqinfo,
+            block_tables=block_tables,
+            page_size=page_size,
+        )
+        paged_bias.k_seqinfo.padding = block_tables.shape[1] * page_size
+        return paged_bias
+
+
+@dataclass
+class PagedBlockDiagonalCausalWithOffsetPaddedKeysMask(AttentionBias):
+    """
+    Same as BlockDiagonalCausalWithOffsetPaddedKeysMask, but for paged attention.
+    block_tables has shape [batch_size, max_num_pages] and K/V have shape
+    [1, max_num_pages * page_size, num_heads, head_dim]
+    or [1, max_num_pages * page_size, num_groups, num_heads, head_dim]
+    """
+
+    q_seqinfo: _SeqLenInfo
+    k_seqinfo: _PaddedSeqLenInfo
+    block_tables: torch.Tensor
+    page_size: int
+
+    def materialize(
+        self,
+        shape: Tuple[int, ...],
+        dtype: torch.dtype = torch.float32,
+        device: Union[str, torch.device] = "cpu",
+    ) -> torch.Tensor:
+        """Materialize the attention bias - for debugging & testing"""
+        # First create a non-paged mask, then cut individual pages and
+        # copy them to their places in the physical mask, using block tables
+
+        max_row_len = self.block_tables.shape[1] * self.page_size
+        logical_input_len = self.block_tables.shape[0] * max_row_len
+        bias_nonpaged = BlockDiagonalCausalWithOffsetPaddedKeysMask(
+            q_seqinfo=self.q_seqinfo,
+            k_seqinfo=_PaddedSeqLenInfo.from_seqlens_padded(
+                self.k_seqinfo.seqlen_py, max_row_len
+            ),
+        )
+        mask_nonpaged = bias_nonpaged.materialize(shape, dtype, device)
+
+        mask_paged = torch.empty(
+            mask_nonpaged.shape[:-1] + (logical_input_len,), dtype=dtype, device=device
+        )
+        mask_paged.fill_(-math.inf)
+        for b, (q_start, q_end) in enumerate(self.q_seqinfo.intervals()):
+            for logical_page_idx in range(self.block_tables.shape[1]):
+                physical_page_idx = self.block_tables[b][logical_page_idx]
+                k_logical_start = (
+                    b * self.block_tables.shape[1] + logical_page_idx
+                ) * self.page_size
+                k_logical_end = k_logical_start + self.page_size
+                k_physical_start, k_physical_end = (
+                    physical_page_idx * self.page_size,
+                    (physical_page_idx + 1) * self.page_size,
+                )
+                mask_paged[
+                    ..., q_start:q_end, k_physical_start:k_physical_end
+                ] = mask_nonpaged[..., q_start:q_end, k_logical_start:k_logical_end]
+
+        return mask_paged
+
+    @classmethod
+    def from_seqlens(
+        cls,
+        q_seqlen: Sequence[int],
+        kv_seqlen: Sequence[int],
+        block_tables: torch.Tensor,
+        page_size: int,
+    ) -> "PagedBlockDiagonalCausalWithOffsetPaddedKeysMask":
+        """Creates a :attr:`BlockDiagonalCausalWithOffsetPaddedKeysMask` from a list of tensor
+        lengths for query and key/value.
+
+        Args:
+            q_seqlen (Sequence[int]): List or tensor of sequence lengths for query tensors
+            kv_padding (int): Padding for k/v - also an upperbound on each individual key length
+            kv_seqlen (Sequence[int]): List or tensor of sequence lengths for key/value.
+            causal_diagonal: unused, for BC only
+        Returns:
+            BlockDiagonalCausalWithOffsetPaddedKeysMask
+        """
+        assert kv_seqlen is None or len(q_seqlen) == len(kv_seqlen), (
+            q_seqlen,
+            kv_seqlen,
+        )
+        q_seqinfo = _SeqLenInfo.from_seqlens(q_seqlen)
+        k_seqinfo = _PaddedSeqLenInfo.from_seqlens_padded(
+            kv_seqlen, padding=block_tables.shape[1] * page_size
+        )
+        return cls(
+            q_seqinfo=q_seqinfo,
+            k_seqinfo=k_seqinfo,
+            block_tables=block_tables,
+            page_size=page_size,
+        )
+
 
 @dataclass
 class BlockDiagonalCausalLocalAttentionMask(BlockDiagonalCausalMask):

{xformers-0.0.25.dev745 → xformers-0.0.25.dev746}/xformers/ops/fmha/common.py

@@ -288,11 +288,12 @@ class AttentionOpBase(BaseOperator):
         Returns a list of reasons why this is not supported.
         The kernel can run these inputs only if the returned list is empty
         """
+        query_shape = d.query.shape
         reasons = cls.shape_not_supported_reasons(
-            Mq=d.query.shape[1],
+            Mq=query_shape[1],
             Mkv=d.key.shape[1],
-            K=d.query.shape[-1],
-            Kv=d.value.shape[-1],
+            K=query_shape[-1],
+            Kv=query_shape[-1] if d.value.dtype == torch.int32 else d.value.shape[-1],
         )
         device_type = d.query.device.type
         dtype = d.query.dtype

{xformers-0.0.25.dev745 → xformers-0.0.25.dev746}/xformers/ops/fmha/flash.py

@@ -135,7 +135,7 @@ try:
                 query,
                 key,
                 value,
-                None,
+                None,  # out
                 cu_seq_lens_q,
                 cu_seq_lens_k,
                 seqused_k,
@@ -286,7 +286,7 @@ def _convert_input_format(
         max_seqlen_q = inp.query.shape[1]
         max_seqlen_k = inp.key.shape[1]
 
-    if query.ndim == 5:  # QGA
+    if query.ndim == 5:  # GQA
         assert supports_mqa
 
         # Fold the group/head_in_group dimensions together