flash_attention_v3_varlen

paddle.nn.functional. flash_attention_v3_varlen ( query: Tensor, key: Tensor, value: Tensor, cu_seqlens_q: Tensor, cu_seqlens_k: Tensor, dropout: float = 0.0, causal: bool = False, return_softmax: Literal[False] = False, *, fixed_seed_offset: Tensor | None = None, rng_name: str = '', training: bool = True, softmax_scale: float | None = None, max_seqlen_q: int = 0, max_seqlen_k: int = 0, name: str | None = None ) → tuple[Tensor, None] [source]
paddle.nn.functional. flash_attention_v3_varlen ( query: Tensor, key: Tensor, value: Tensor, cu_seqlens_q: Tensor, cu_seqlens_k: Tensor, dropout: float = 0.0, causal: bool = False, return_softmax: Literal[True] = False, *, fixed_seed_offset: Tensor | None = None, rng_name: str = '', training: bool = True, softmax_scale: float | None = None, max_seqlen_q: int = 0, max_seqlen_k: int = 0, name: str | None = None ) → tuple[Tensor, Tensor]

The equation is:

\[result=softmax(\frac{ Q * K^T }{\sqrt{d}}) * V\]

where : Q, K, and V represent the three input parameters of the attention module. The dimensions of the three parameters are the same. d represents the size of the last dimension of the three parameters. This is the varlen version of flash attention.

Warning

This API is only support inputs with dtype float16 and bfloat16.

Parameters

  • query (Tensor) – The query tensor in the Attention module. 3-D tensor with shape: [token_num, num_heads, head_dim]. The dtype can be float16 or bfloat16.

  • key (Tensor) – The key tensor in the Attention module. 3-D tensor with shape: [token_num, num_heads, head_dim]. The dtype can be float16 or bfloat16.

  • value (Tensor) – The value tensor in the Attention module. 3-D tensor with shape: [token_num, num_heads, head_dim]. The dtype can be float16 or bfloat16.

  • cu_seqlens_q (Tensor) – The cumulative sequence lengths of the sequences in the batch, used to index query.

  • cu_seqlens_k (Tensor) – The cumulative sequence lengths of the sequences in the batch, used to index key and value.

  • dropout (float) – The dropout ratio.

  • causal (bool) – Whether enable causal mode.

  • return_softmax (bool) – Whether to return softmax.

  • fixed_seed_offset (Tensor|None, optional) – With fixed seed, offset for dropout mask.

  • rng_name (str) – The name to select Generator.

  • training (bool) – Whether it is in the training phase.

  • softmax_scale (float) – The softmax scale of the attention.

  • max_seqlen_q (int) – Maximum sequence length of query in the batch. Note it’s the padding length, not the max actual seqlen.

  • max_seqlen_k (int) – Maximum sequence length of key/value in the batch.

  • name (str|None, optional) – The default value is None. Normally there is no need for user to set this property. For more information, please refer to Name.

Returns

The attention tensor. 3-D tensor with shape: [token_num, num_heads, head_dim]. The dtype can be float16 or bfloat16. softmax(Tensor): The softmax tensor. None if return_softmax is False.

Return type

out(Tensor)

Examples

>>> 
>>> import paddle

>>> paddle.seed(2023)
>>> q = paddle.rand((10, 2, 128), dtype="bfloat16")
>>> cu_seqlens_q = paddle.to_tensor([0, 10], dtype="int32")
>>> max_seq_len_q = 10

>>> output = paddle.nn.functional.flash_attention.flash_attention_v3_varlen(q, q, q, cu_seqlens_q, cu_seqlens_q, max_seqlen_q=max_seq_len_q, max_seqlen_k=max_seq_len_q, causal=True)
>>>