Function attention_python

python/tvm/topi/testing/attention_python.py:25–125

Attention operator in python Parameters ---------- q : np.ndarray Query tensor with shape [batch, seq_length, num_heads, head_dim] in the layout specified by `layout`. k : np.ndarray Key tensor with shape [batch, seq_length_kv, num_kv_heads, head_dim] in the

(
    q: np.ndarray,
    k: np.ndarray,
    v: np.ndarray,
    bias: np.ndarray | None,
    qk_scale: float,
    causal: str,
    window_size: int | None = None,
    layout: str = "BSNH",
)

def attention_python(
    q: np.ndarray,
    k: np.ndarray,
    v: np.ndarray,
    bias: np.ndarray | None,
    qk_scale: float,
    causal: str,
    window_size: int | None = None,
    layout: str = "BSNH",
):  # pylint: disable=too-many-arguments, too-many-locals, invalid-name
    """Attention operator in python

    Parameters
    ----------
    q : np.ndarray
        Query tensor with shape [batch, seq_length, num_heads, head_dim] in the layout specified by
        `layout`.
    k : np.ndarray
        Key tensor with shape [batch, seq_length_kv, num_kv_heads, head_dim] in the layout specified
        by `layout`.
    v : np.ndarray
        Value tensor with shape [batch, seq_length_kv, num_kv_heads, head_dim_v] in the layout
        specified by `layout`.
    bias : np.ndarray
        Bias tensor with shape [batch, num_heads, seq_length, seq_length]
    qk_scale : float
        Scale factor for the query-key product.
    causal : str
        The type of causal mask to apply. Can be "none", "TopLeft", or "BottomRight".
    window_size : Optional[int]
        The window size for the causal mask.
    layout : str
        The layout of the input tensors, e.g. "BSNH" or "BNSH".

    Returns
    -------
    np.ndarray
        The output tensor with shape [batch, seq_length, num_heads, head_dim_v] in the layout
        specified by `layout`.
    """
    assert layout in ["BSNH", "BNSH", "SBNH"]

    dim_b = layout.find("B")
    dim_s = layout.find("S")
    dim_n = layout.find("N")
    dim_h = layout.find("H")

    q = q.transpose(dim_b, dim_n, dim_s, dim_h)  # b, n, s, h
    k = k.transpose(dim_b, dim_n, dim_s, dim_h)  # b, n, s_kv, h
    kt = k.transpose(0, 1, 3, 2)  # b, n, h, s_kv
    v = v.transpose(dim_b, dim_n, dim_s, dim_h)

    num_heads = q.shape[1]
    num_kv_heads = k.shape[1]
    s = q.shape[2]
    s_kv = k.shape[2]

    if num_heads != num_kv_heads: