xformers 0.0.22.dev605__tar.gz → 0.0.22.dev606__tar.gz

{xformers-0.0.22.dev605/xformers.egg-info → xformers-0.0.22.dev606}/PKG-INFO

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

xformers-0.0.22.dev606/tests/test_rope_padded.py

@@ -0,0 +1,168 @@
+# 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 Optional
+
+import pytest
+import torch
+
+from xformers.ops import rope_padded
+from xformers.ops.fmha.attn_bias import BlockDiagonalCausalWithOffsetPaddedKeysMask
+
+from .utils import assert_allclose
+
+compute_capability = (0, 0)
+if torch.cuda.is_available():
+    compute_capability = torch.cuda.get_device_capability("cuda")
+cuda_sm80_only = pytest.mark.skipif(
+    compute_capability < (8, 0), reason="requires sm80+"
+)
+
+
+def _slow_rope(
+    x: torch.Tensor,
+    *,
+    seqpos: Optional[torch.Tensor] = None,
+    theta=10000,
+    adjacents: bool = True,
+):
+    """
+    Simple rope calculation of rope of one tensor
+
+    Args:
+        x: input, shape (B, M, H, K).
+        seqpos: gives the position of each sequence element in x in its sequence
+            (shape (M,)).
+    """
+    B, M, H, dim = x.shape
+    assert dim % 2 == 0
+    if seqpos is None:
+        seqpos = torch.arange(M, device=x.device)
+    power = torch.arange(0, dim, 2, device=x.device)[: (dim // 2)].float() / dim
+    freqs = 1.0 / (theta**power)
+    all_freqs = torch.outer(seqpos, freqs)
+    freqs_cis = torch.polar(torch.ones_like(all_freqs), all_freqs)  # complex64
+    if adjacents:
+        x_reshaped = x.float().reshape(B, M, H, -1, 2)
+        x_ = torch.view_as_complex(x_reshaped)
+        x_out = torch.view_as_real(x_ * freqs_cis[:, None])
+    else:
+        x_reshaped = x.float().reshape(B, M, H, 2, -1).transpose(3, 4).contiguous()
+        x_ = torch.view_as_complex(x_reshaped)
+        x_out = torch.view_as_real(x_ * freqs_cis[:, None])
+        x_out = x_out.transpose(3, 4)
+    return x_out.flatten(3).type_as(x)
+
+
+def _slow_rope2(
+    x: torch.Tensor,
+    *,
+    seqpos: Optional[torch.Tensor] = None,
+    theta=10000,
+    adjacents: bool = True,
+):
+    """
+    More flexible unused version of _slow_rope
+    - allows varying dtypes.
+    """
+    internal_dtype = torch.float64
+    B, M, H, dim = x.shape
+    assert dim % 2 == 0
+    if seqpos is None:
+        seqpos = torch.arange(M, device=x.device)
+    power = (
+        torch.arange(0, dim, 2, device=x.device)[: (dim // 2)].to(internal_dtype) / dim
+    )
+    # freqs = 1.0 / (theta**power)
+    freqs = theta**-power
+    f = torch.outer(seqpos, freqs)[:, None]
+    if adjacents:
+        x1, x2 = x.to(internal_dtype).reshape(B, M, H, -1, 2).unbind(-1)
+        y1 = x1 * f.cos() - x2 * f.sin()
+        y2 = x1 * f.sin() + x2 * f.cos()
+        x_out = torch.stack([y1, y2], -1)
+    else:
+        x1, x2 = x.to(internal_dtype).reshape(B, M, H, 2, -1).unbind(3)
+        y1 = x1 * f.cos() - x2 * f.sin()
+        y2 = x1 * f.sin() + x2 * f.cos()
+        x_out = torch.stack([y1, y2], -2)
+    return x_out.flatten(3).type_as(x)
+
+
+DTYPES = {"bf16": torch.bfloat16, "f32": torch.float32}
+
+
+@cuda_sm80_only
+@pytest.mark.parametrize(
+    "adjacents", [True, False], ids=lambda x: "adj" if x else "non-adj"
+)
+@pytest.mark.parametrize("dtype_str", ["bf16", "f32"])
+@pytest.mark.parametrize("internal_dtype", ["", "f32", "f64"])
+@pytest.mark.parametrize("dim", [100, 4098])
+@pytest.mark.parametrize("padding", [87, 18300])
+def test_consistency(
+    adjacents: bool, dim: int, padding: int, internal_dtype: str, dtype_str: str
+):
+    torch.manual_seed(1)
+    heads, kvheads = 10, 2
+    nqueries = [2, 1, 1]
+    cache_lens = [27, padding - 5, padding // 2]
+    device = torch.device("cuda")
+    dtype = DTYPES[dtype_str]
+
+    # Can we make the internals of attn_bias be on the gpu.
+    attn_bias = BlockDiagonalCausalWithOffsetPaddedKeysMask.from_seqlens(
+        q_seqlen=nqueries, kv_padding=padding, kv_seqlen=cache_lens
+    )
+
+    total_cache_length = len(cache_lens) * padding
+    total_nqueries = sum(nqueries)
+    cache_k = torch.rand(
+        1, total_cache_length, kvheads, dim, device=device, dtype=dtype
+    )
+    cache_v = torch.rand(
+        1, total_cache_length, kvheads, dim, device=device, dtype=dtype
+    )
+    cache_k_orig = cache_k.clone()
+    cache_v_orig = cache_v.clone()
+    xq = torch.rand(1, total_nqueries, heads, dim, device=device, dtype=dtype)
+    xk = torch.rand(1, total_nqueries, kvheads, dim, device=device, dtype=dtype)
+    xv = torch.rand(1, total_nqueries, kvheads, dim, device=device, dtype=dtype)
+
+    out = rope_padded(
+        xq,
+        xk,
+        xv,
+        cache_k,
+        cache_v,
+        attn_bias,
+        adjacents=adjacents,
+        internal_dtype=internal_dtype,
+    )
+
+    seqpos = torch.tensor(
+        [cache_lens[0] - 2, cache_lens[0] - 1, cache_lens[1] - 1, cache_lens[2] - 1],
+        device=device,
+    )
+    cache_locs = [seqpos[0], seqpos[1], padding + seqpos[2], 2 * padding + seqpos[3]]
+    baseline = _slow_rope if dtype_str == "f32" else _slow_rope2
+    expected_out = baseline(xq, seqpos=seqpos, adjacents=adjacents)
+    atol = 5e-3 if dim > 4000 and dtype_str != "f32" else 3e-3
+    rtol = 8e-3 if dtype_str == "bf16" else 1e-5
+    assert_allclose(out, expected_out, atol=atol, rtol=rtol)
+
+    assert_allclose(cache_v[:, cache_locs], xv, atol=atol, rtol=rtol)
+    cache_v[:, cache_locs] = cache_v_orig[:, cache_locs]
+    assert torch.allclose(cache_v, cache_v_orig)
+
+    slow_roped_xk = _slow_rope(xk, seqpos=seqpos, adjacents=adjacents)
+    assert_allclose(
+        cache_k[:, cache_locs],
+        slow_roped_xk,
+        atol=atol,
+        rtol=rtol,
+    )
+    cache_k[:, cache_locs] = cache_k_orig[:, cache_locs]
+    assert torch.allclose(cache_k, cache_k_orig)

{xformers-0.0.22.dev605 → xformers-0.0.22.dev606}/xformers/ops/__init__.py

@@ -24,6 +24,7 @@ from .fmha import (
 )
 from .indexing import index_select_cat, scaled_index_add
 from .rmsnorm import RMSNorm
+from .rope_padded import rope_padded
 from .swiglu_op import (
     SwiGLU,
     SwiGLUEagerOp,
@@ -92,4 +93,5 @@ __all__ = [
     "masked_matmul",
     "scaled_index_add",
     "index_select_cat",
+    "rope_padded",
 ]

xformers-0.0.22.dev606/xformers/ops/rope_padded.py

@@ -0,0 +1,188 @@
+# 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 Optional, Tuple
+
+import torch
+
+from xformers.ops.fmha.attn_bias import (  # type: ignore
+    BlockDiagonalCausalWithOffsetPaddedKeysMask,
+)
+
+from .. import _is_triton_available
+
+
+def rope_padded(
+    xq: torch.Tensor,
+    xk: torch.Tensor,
+    xv: torch.Tensor,
+    cache_k: torch.Tensor,
+    cache_v: torch.Tensor,
+    attn_bias: BlockDiagonalCausalWithOffsetPaddedKeysMask,
+    *,
+    theta: float = 10000.0,
+    out_q: Optional[torch.Tensor] = None,
+    adjacents: bool = True,
+    internal_dtype: str = "",
+):
+    """
+    Performs RoPE (rotary embeddings) and kv-cache emplacement for a heterogeneous
+    batch for inference in the style given by
+    BlockDiagonalCausalWithOffsetPaddedKeysMask.
+    The batch is concatted along the sequence dimension, so the
+    actual dim-0 length of all tensors is 1.
+
+    xq, xk and xv should be (1, slen, n_heads, dim), where
+    xq's n_heads can differ from xk and xv.
+
+    This function places the roped xk in the right place in cache_k, and
+    xv (unmodified) in the right place in cache_v, and returns out_q
+    (the roped xq) such that things are ready to call
+
+    xformers.ops.memory_efficient_attention(
+        out_q, cache_k, cache_v, attn_bias=attn_bias
+    )
+
+    This functionality is experimental. Its API might be changed without warnings.
+    Use it at your own risk.
+
+    Arguments:
+        xq: tensor of queries to apply rope to
+        xk: tensor of keys to apply rope to
+        xv: tensor of values to copy into cache_v
+        cache_k: cache of keys, MODIFIED IN PLACE
+        cache_v: cache of values, MODIFIED IN PLACE
+        attn_bias: details the layout of caches.
+                Used to determine frequencies for the
+                RoPE calculation as well as the locations in cache_k and cache_v
+                to write to. Must be on the device.
+        adjacents: If True, the inputs are in adjacent pairs along the final dim axis.
+                  This is like the released LLaMA model.
+                  If False, the dim axis is split in two equal pieces.
+                   I.e. the features are ordered with all the real parts before all
+                   the imaginary parts. This matches HuggingFace, e.g.
+                   https://github.com/huggingface/transformers/blob/
+                   f143037789288ba532dada934a118e648e715738/
+                   src/transformers/models/llama/modeling_llama.py#L126-L130
+        internal_dtype: set to "f32" or "f64" to enforce dtype in the calculation
+    """
+    if torch.is_grad_enabled() and (
+        xq.requires_grad
+        or xk.requires_grad
+        or xv.requires_grad
+        or cache_k.requires_grad
+        or cache_v.requires_grad
+        or out_q is not None
+    ):
+        raise ValueError("Gradients not supported.")
+    assert _is_triton_available()
+    import triton
+
+    from .triton.rope_padded_kernels import _rope_padded_kernel
+
+    n_total_queries = attn_bias.q_seqinfo.seqstart_py[-1]
+    cache_length = attn_bias.k_seqinfo.seqstart_py[-1]
+    bsz, q_len, n_q_heads, dim = xq.shape
+    assert q_len == n_total_queries
+    if bsz != 1:
+        raise ValueError(
+            "Expected batch size dimension to be 1" "as batches should be concatenated."
+        )
+    xk_shape = xk.shape
+    n_kv_heads = xk_shape[2]
+    if xk_shape != (1, n_total_queries, n_kv_heads, dim):
+        raise ValueError("unexpected k shape")
+    if xv.shape != (1, n_total_queries, n_kv_heads, dim):
+        raise ValueError("unexpected v shape")
+    if cache_k.shape != (1, cache_length, n_kv_heads, dim):
+        raise ValueError("unexpected cache_k length")
+    if cache_v.shape != (1, cache_length, n_kv_heads, dim):
+        raise ValueError("unexpected cache_v length")
+
+    xq_stride = xq.stride()
+    xk_stride = xk.stride()
+    xv_stride = xv.stride()
+    cache_k_stride = cache_k.stride()
+    cache_v_stride = cache_v.stride()
+    if xq_stride[3] != 1:
+        raise ValueError("Each q head must be contiguous")
+    if xk_stride[3] != 1:
+        raise ValueError("Each k head must be contiguous")
+    if xv_stride[3] != 1:
+        raise ValueError("Each v head must be contiguous")
+    if cache_k_stride[3] != 1:
+        raise ValueError("Each cache_k head must be contiguous")
+    if cache_v_stride[3] != 1:
+        raise ValueError("Each cache_v head must be contiguous")
+    n_total_heads = n_q_heads + 2 * n_kv_heads
+    v_start = n_total_heads - n_kv_heads
+    k_start = n_q_heads
+    if out_q is None:
+        out_q = xq.new_empty(1, n_total_queries, n_q_heads, dim)
+        out_q_stride: Tuple[int, ...] = (0, n_q_heads * dim, dim, 1)
+    else:
+        if out_q.shape != xq.shape:
+            raise ValueError("Unexpected shape of out_q")
+        out_q_stride = out_q.stride()
+        if out_q_stride[3] != 1:
+            raise ValueError("Each out_q head must be contiguous")
+
+    assert out_q is not None
+
+    logical_bsz = len(attn_bias.q_seqinfo.seqstart_py) - 1
+
+    # Less than 64KB per feature: enqueue fused kernel
+    MAX_FUSED_SIZE = 65536 // xq.element_size()
+    BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(dim))
+    BLOCK_SIZE = max(BLOCK_SIZE, 128)
+    BLOCK_SIZE = min(BLOCK_SIZE, 4096)
+    # heuristics for number of warps
+    num_warps = min(max(BLOCK_SIZE // 256, 1), 8)
+    device = xq.device
+    # Move these to the right device, like fmha does.
+    attn_bias.k_seqinfo.to(device)
+    attn_bias.q_seqinfo.to(device)
+    seqstartq = attn_bias.q_seqinfo.seqstart
+    seqstartk = attn_bias.k_seqinfo.seqstart
+    seqlenk = attn_bias.k_seqinfo.seqlen
+    assert internal_dtype in ["", "f32", "f64"]
+    # experiment with the order of dims here.
+    _rope_padded_kernel[(logical_bsz, attn_bias.q_seqinfo.max_seqlen, n_total_heads)](
+        xq,
+        xk,
+        xv,
+        out_q,
+        cache_k,
+        cache_v,
+        seqstartq,
+        seqstartk,
+        seqlenk,
+        theta,
+        k_start,
+        v_start,
+        dim,
+        xq_stride[1],
+        xq_stride[2],
+        xk_stride[1],
+        xk_stride[2],
+        xv_stride[1],
+        xv_stride[2],
+        cache_k_stride[1],
+        cache_k_stride[2],
+        cache_v_stride[1],
+        cache_v_stride[2],
+        seqstartq.stride(0),
+        seqstartk.stride(0),
+        seqlenk.stride(0),
+        out_q_stride[1],
+        out_q_stride[2],
+        internal_dtype,
+        const_batch_strides=False,
+        cache_padding_length=0,
+        seqlenk_shift=0,
+        BLOCK_SIZE=BLOCK_SIZE,
+        adjacents=adjacents,
+        num_warps=num_warps,
+    )
+    return out_q

xformers-0.0.22.dev606/xformers/ops/triton/rope_padded_kernels.py

@@ -0,0 +1,161 @@
+# 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.
+import triton  # type: ignore
+import triton.language as tl  # type: ignore
+
+if hasattr(tl, "libdevice"):
+    tl_math = tl.libdevice
+else:
+    tl_math = tl.math
+
+
+@triton.jit
+def _rope_padded_kernel(
+    xq,
+    xk,
+    xv,
+    out_q,
+    cache_k,
+    cache_v,
+    seqstartq,
+    seqstartk,
+    seqlenk,
+    theta,
+    k_start: tl.constexpr,
+    v_start: tl.constexpr,
+    dim: tl.constexpr,  # dimension of each head
+    stride_xqM,
+    stride_xqH,
+    stride_xkM,
+    stride_xkH,
+    stride_xvM,
+    stride_xvH,
+    stride_cachekM,
+    stride_cachekH,
+    stride_cachevM,
+    stride_cachevH,
+    stride_seqstartq,
+    stride_seqstartk,
+    stride_seqlenk,
+    stride_outqM,
+    stride_outqH,
+    internal_dtype: tl.constexpr,
+    # If True, seqstartq and seqstartk are not used but rather we
+    # assume that every batch element has the same number of
+    # queries (i.e. num_queries := tl.num_programs(1) )
+    # and the same cache space cache_padding_length.
+    # Always False when called below.
+    const_batch_strides: tl.constexpr,
+    # If const_batch_strides==True, the common cache length for each batch element.
+    # (Only the first seqlenk[i] elements are actually in use, and only the last
+    #  num_queries of those are actually written to.)
+    cache_padding_length,
+    # offset added to all values in seqlenk before using them.
+    # Always 0 when called below.
+    seqlenk_shift: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    adjacents: tl.constexpr,
+):
+    """
+    Each letter in this diagram is a whole row of length dim.
+
+     INPUT      xq        xk       xv
+
+        head_dim ─►
+
+      batch   qqqqqq      kk       vv
+        │     qqqqqq      kk       vv
+        ▼     qqqqqq      kk       vv
+
+    head_idx:  (goes across all heads of all 3 inputs)
+              ▲     ▲     ▲ ▲      ▲ ▲
+              │     │     │ │      │ │
+                          │        │
+              0  k_start  │v_start │n_total_heads
+                          │        │
+                          │        │
+                      k_start    v_start
+
+    Output is to out_q (same shape as xq), an xk-shaped part
+    of cache_k and an xv-shaped part of cache_v
+    """
+    batch_elt = tl.program_id(0)
+    query_pos_in_batch_elt = tl.program_id(1)
+    head_idx = tl.program_id(2)
+
+    if internal_dtype == "f32":
+        theta = theta.to(tl.float32)
+    elif internal_dtype == "f64":
+        theta = theta.to(tl.float64)
+
+    if const_batch_strides:
+        query_pos = query_pos_in_batch_elt + tl.num_programs(1) * batch_elt
+        end_query_pos = tl.num_programs(1) * (batch_elt + 1)
+    else:
+        query_pos = query_pos_in_batch_elt + tl.load(
+            seqstartq + batch_elt * stride_seqstartq
+        )
+        end_query_pos = tl.load(seqstartq + (batch_elt + 1) * stride_seqstartq)
+        if query_pos >= end_query_pos:
+            return
+
+    is_q = head_idx < k_start
+    is_v = head_idx >= v_start
+
+    xq += query_pos * stride_xqM + head_idx * stride_xqH
+    out_q += query_pos * stride_outqM + head_idx * stride_outqH
+
+    if const_batch_strides:
+        cache_start = cache_padding_length * batch_elt
+    else:
+        cache_start = tl.load(seqstartk + batch_elt * stride_seqstartk)
+    end_of_batch_elt_cache = (
+        cache_start + tl.load(seqlenk + batch_elt * stride_seqlenk) + seqlenk_shift
+    )
+
+    cache_pos = end_of_batch_elt_cache - (end_query_pos - query_pos)
+    seq_pos = cache_pos - cache_start
+    cache_k += (head_idx - k_start) * stride_cachekH + cache_pos * stride_cachekM
+    xk += query_pos * stride_xkM + (head_idx - k_start) * stride_xkH
+    in_qk = tl.where(is_q, xq, xk)
+    out_qk = tl.where(is_q, out_q, cache_k)
+
+    cache_v += (head_idx - v_start) * stride_cachevH + cache_pos * stride_cachevM
+    xv += query_pos * stride_xvM + (head_idx - v_start) * stride_xvH
+
+    out = tl.where(is_v, cache_v, out_qk)
+    x_in = tl.where(is_v, xv, in_qk)
+
+    for offset in range(0, dim // 2, BLOCK_SIZE // 2):
+        c = tl.arange(0, BLOCK_SIZE // 2)
+        powers = (offset + c) * 2.0
+        if adjacents:
+            cols_re = (offset + c) * 2
+            cols_im = cols_re + 1
+        else:
+            cols_re = offset + c
+            cols_im = cols_re + dim // 2
+
+        mask = cols_im < dim
+
+        re_x = tl.load(x_in + cols_re, mask=mask)
+        im_x = tl.load(x_in + cols_im, mask=mask)
+        # freqs = seq_pos / (theta ** (powers / dim))
+        freqs = seq_pos * tl_math.pow(theta, powers / (-dim))
+        sines = tl.sin(freqs)
+        cosines = tl.cos(freqs)
+        re_out = re_x * cosines - im_x * sines
+        im_out = im_x * cosines + re_x * sines
+
+        re_out_ = tl.where(is_v, re_x, re_out)
+        im_out_ = tl.where(is_v, im_x, im_out)
+        if internal_dtype == "f64":
+            if re_x.dtype == tl.bfloat16:
+                # triton 2.0.0 crashes if you try to convert
+                # float64 directly to bfloat16, so make an intermediate step.
+                re_out_ = re_out_.to(tl.float32)
+                im_out_ = im_out_.to(tl.float32)
+        tl.store(out + cols_re, re_out_, mask=mask)
+        tl.store(out + cols_im, im_out_, mask=mask)

xformers-0.0.22.dev606/xformers/version.py

@@ -0,0 +1,2 @@
+# noqa: C801
+__version__ = "0.0.22.dev606"

{xformers-0.0.22.dev605 → xformers-0.0.22.dev606/xformers.egg-info}/PKG-INFO

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

{xformers-0.0.22.dev605 → xformers-0.0.22.dev606}/xformers.egg-info/SOURCES.txt

@@ -59,6 +59,7 @@ tests/test_profiler.py
 tests/test_residual.py
 tests/test_reversible.py
 tests/test_rmsnorm.py
+tests/test_rope_padded.py
 tests/test_rotary_embeddings.py
 tests/test_sparse_tensors.py
 tests/test_sparsecs.py
@@ -2024,6 +2025,7 @@ xformers/ops/__init__.py
 xformers/ops/common.py
 xformers/ops/indexing.py
 xformers/ops/rmsnorm.py
+xformers/ops/rope_padded.py
 xformers/ops/swiglu_op.py
 xformers/ops/unbind.py
 xformers/ops/fmha/__init__.py
@@ -2037,6 +2039,7 @@ xformers/ops/fmha/small_k.py
 xformers/ops/fmha/triton.py
 xformers/ops/triton/__init__.py
 xformers/ops/triton/rmsnorm_kernels.py
+xformers/ops/triton/rope_padded_kernels.py
 xformers/profiler/__init__.py
 xformers/profiler/api.py
 xformers/profiler/device_limits.py

xformers-0.0.22.dev605/xformers/version.py

@@ -1,2 +0,0 @@
-# noqa: C801
-__version__ = "0.0.22.dev605"