flash_attention_with_sparse_mask

paddle.nn.functional. flash_attention_with_sparse_mask ( query, key, value, attn_mask_start_row_indices, attn_mask_start_row=0, dropout_p=0.0, is_causal=False, return_softmax=False, return_softmax_lse=False, return_seed_offset=False, training=True, name=None ) [源代码]

用稀疏 mask 表达的 flash_attention。

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

参数

  • query (int) - 输入 Query Tensor，shape =[batch_size, seq_len, num_heads, head_dim]，数据类型为 float16 或 bfloat16。

  • key (Tensor) - 输入 Key Tensor，shape 以及 dtype 和 query 相同。

  • value (Tensor) - 输入 Value Tensor，shape 以及 dtype 和 query 相同。

  • attn_mask_start_row_indices (Tensor) - 稀疏掩码索引，shape =[batch_size, num_head, seq_len]，每个元素的值表示得分矩阵中掩码开始的行索引。数据类型必须是 int32。

  • attn_mask_start_row (Tensor，可选) - 当传入 attn_mask_start_row_indices 并且已知最小行数大于 0 时，可以设置 attn_mask_start_row 以提高性能。默认值为 0。

  • dropout_p (bool，可选) – dropout 概率值，默认值为 0。

  • is_causal (bool，可选) - 是否使用 causal 模式，默认值：False。

  • return_softmax (bool，可选) - 是否返回 softmax 的结果。默认值 False。

  • return_softmax_lse (bool，可选) - 是否返回 return_softmax_lse 的结果。默认值为 False。

  • return_seed_offset (bool，可选) - 是否返回 return_seed_offset 的结果。默认值为 False。

  • training (bool，可选) - 指示是否为训练模式。默认值为 True。

  • name (str，可选) - 具体用法请参见 Name，一般无需设置，默认值为 None。

返回

Tensor，attention 的结果。

代码示例

>>> import paddle
>>> import numpy as np
>>> def generate_start_rows(bz, num_head, rows, cols, start_row):
>>>     assert rows == cols, f"rows {rows} must be equal to cols {cols}."
>>>     start_rows_list = []
>>>     for bz_idx in range(bz):
>>>         for head_idx in range(num_head):
>>>             start_rows = np.array([rows+1] * cols)
>>>             mask_pos = np.random.choice(cols-1, cols - start_row, replace=False)
>>>             index = np.arange(start_row, rows)
>>>             mask_pos = np.concatenate([mask_pos[mask_pos < index - 1], mask_pos[mask_pos >= index - 1]])
>>>             start_rows[mask_pos] = index
>>>             start_rows_list.append(start_rows)
>>>     start_rows_arr = np.array(start_rows_list).reshape([bz, num_head, rows])
>>>     return start_rows_arr
>>> q = paddle.rand((1, 128, 2, 16), dtype=paddle.bfloat16)
>>> attn_mask_start_row = 48
>>> start_row_indices = generate_start_rows(1, 2, 128, 128, attn_mask_start_row)
>>> attn_mask_start_row_indices = paddle.to_tensor(start_row_indices, dtype=paddle.int32)
>>> out = paddle.nn.functional.flash_attention.flash_attention_with_sparse_mask(
>>>     q, q, q,
>>>     attn_mask_start_row_indices=attn_mask_start_row_indices,
>>>     attn_mask_start_row=attn_mask_start_row,
>>>     dropout_p=0.9,
>>>     is_causal=True,
>>> )
>>> print(out)