flash_attn

Applicable Products

• Ascend 950PR/Ascend 950DT: Supported
• Atlas A3 training products/Atlas A3 inference products: Not Supported
• Atlas A2 training products/Atlas A2 inference products: Not Supported
• Atlas 200I/500 A2 inference products: Not Supported
• Atlas inference products: Not Supported
• Atlas training products: Not Supported

Function

• Description:

  flash_attn is an extension API of cann_ops_transformer based on TorchNPU. It calls the FlashAttn operator to perform non-quantized attention computation with shared KV (Key and Value use the same input), unified for training and inference.

  flash_attn_metadata is the metadata generation API of flash_attn. It generates metadata before the main operator executes. The metadata records the task splitting results of AICore/AIVCore. The main operator can optionally take this metadata to optimize scheduling. The typical call flow is as follows:

  1. Prepare inputs such as q, k, and v.
  2. Call flash_attn_metadata to generate metadata.
  3. Call flash_attn and pass the metadata obtained in the previous step to the main operator.

• Formula:

  Self-attention builds an attention model using the relationships within the input samples themselves. The principle is as follows: assume an input sample sequence $x$ of length $n$, where each element of $x$ is a $d$-dimensional vector. Each $d$-dimensional vector can be regarded as a token embedding. Transform this sequence through three weight matrices to obtain three matrices of dimension $n \times d$.

  The self-attention formula is generally defined as follows, where $Q$, $K$, and $V$ are important property elements of the input sample, obtained by spatially transforming the input sample and unifiable into a single feature space. "Attention" in the formula and operator name is an abbreviation of "self-attention".

  $$Attention(Q,K,V)=Score(Q,K)V$$

  In this operator, the Score function uses the Softmax function. The self-attention formula is:

  $$Attention(Q,K,V)=Softmax(\frac{QK^T}{\sqrt{d}})V$$

  Where the product of $Q$ and $K^T$ represents the attention of the input $x$. To prevent this value from becoming too large, it is usually divided by $\sqrt{d}$ for scaling, and each row is normalized by softmax. After multiplying by $V$, a matrix of dimension $n \times d$ is obtained.

  After adding sink, the computation logic is as follows, mainly modifying the softmax_max and softmax_sum logic computation parts.

  $$S = \frac{QK^T}{\sqrt{d}}$$ $$m = max(sink, max(S))$$ $$Attention = \frac{e^{S - m} \times V}{\sum e^{S-m} + e^{sink - m}}$$

After return_softmax_lse is enabled, the returned softmax_lse is calculated as follows:

$$S = \frac{QK^T}{\sqrt{d}}$$ $$softmax\_max = max(S)$$ $$softmax\_lse = log{\sum e^{S-softmax\_max}} + softmax\_max$$

Prototype

Before calling the flash_attn API, call the preceding flash_attn_metadata API to complete the load balancing computation for flash_attn.

flash_attn_metadata(
    num_heads_q,
    num_heads_kv,
    head_dim,
    *,
    cu_seqlens_q=None,
    cu_seqlens_kv=None,
    seqused_q=None,
    seqused_kv=None,
    batch_size=None,
    max_seqlen_q=None,
    max_seqlen_kv=None,
    mask_mode=None,
    win_left=None,
    win_right=None,
    layout_q=None,
    layout_kv=None,
    layout_out=None
) -> Tensor

flash_attn(
    q,
    k,
    v,
    block_table=None,
    cu_seqlens_q=None,
    cu_seqlens_kv=None,
    seqused_q=None,
    seqused_kv=None,
    sinks=None,
    metadata=None,
    softmax_scale=1.0,
    mask_mode=0,
    attn_mask=None,
    win_left=-1,
    win_right=-1,
    max_seqlen_q=-1,
    max_seqlen_kv=-1,
    layout_q="BSND",
    layout_kv="BSND",
    layout_out="BSND",
    return_softmax_lse=False
) -> (Tensor, Tensor)

Parameters

flash_attn_metadata

flash_attn

Return Value

flash_attn_metadata

flash_attn

Note

• attn_out: Tensor type, the output in the formula. Supported data types: float16 and bfloat16. Supported data format: ND. Constraint: the D dimension of this output parameter must be consistent with the D dimension of value, and the other dimensions must be consistent with the shape of the input query.
• softmax_lse: Tensor type. In the ring attention algorithm, the result of query multiplied by key first takes max to obtain softmax_max. The result of query multiplied by key subtracts softmax_max, then takes exp, and finally takes sum to obtain softmax_sum. Finally, log is taken on softmax_sum, and softmax_max is added to obtain the result. Supported data type: float32. When return_softmax_lse is True, the output is generally a Tensor of shape (B, Q_N, Q_S). When input_q is TND, the output is a Tensor of shape (Q_N, Q_T). When return_softmax_lse is False, the output is a Tensor of shape [1] with value 0.

Constraints

Declaration:

• The parameters cu_seqlens_q, cu_seqlens_kv, seqused_q, seqused_kv, block_table, and attn_mask are tensors. Because the operator cannot obtain the specific values of tensors during the Tiling phase, the Tiling side does not validate the values. The user must ensure correctness. If invalid values are passed for the above parameters, undefined behavior is triggered (precision issues, program crashes caused by illegal memory access, etc.).
• The input parameters of flash_attn_metadata and flash_attn must be consistent when called. Because the operator is called in two phases through two APIs, the operator cannot validate this by itself. The user must ensure correctness. If inconsistent parameters are passed to the APIs, undefined behavior occurs (precision issues, program crashes caused by illegal memory access, etc.).

Feature Parameter Group

Basic Information

In the documentation constraints, common fields are defined as follows:

Parameter Group Constraints

Common Parameter Group

• Handling of empty input parameters:

  • An empty tensor means that the shape size of the required input and output is 0, that is, any axis is 0.
  • Use cases that trigger an empty tensor are all intercepted and reported as errors.

• q, k, v, attn_out validation

Parameter	Single-Parameter Validation	Existence Validation	Consistency Validation	Feature Cross-Validation
q	tensor_type supports BFLOAT16 and FLOAT16. BNSD -> (B, N_Q, S_Q, D) BSND -> (B, S_Q, N_Q, D) TND -> (T_Q, N_Q, D)	Must exist.	The data types of q, k, v, and attn_out must be the same. The shapes of k and v must be the same. For layout validation rules, see the layout matching relationship table.	Axis validation: 65536 > B > 0 Q_T > 0 KV_T > 0 Q_N > 0 KV_N > 0 Q_S > 0 KV_S > 0 D supports 64, 128, and 256. The D dimension of q must be equal to that of k/v. Q_N % KV_N == 0 and Q_N / KV_N > 0
k	tensor_type supports BFLOAT16 and FLOAT16. BNSD -> (B, N_KV, S_KV, D) BSND -> (B, S_KV, N_KV, D) TND -> (T_KV, N_KV, D) PA_BBND -> (num_blocks, block_size, KV_N, D) PA_BNBD -> (num_blocks, KV_N, block_size, D) PA_NZ -> (num_blocks, KV_N, D//16, block_size, 16) 1024 >= block_size >= 16, block_size % 16 == 0 Non-contiguous is not supported.
v
attn_out	data_type supports BFLOAT16 and FLOAT16. The shape dimension supports 3 and 4.

Layout matching relationship table:

layout_q	layout_kv	layout_out	layout_softmax_lse
BNSD	BNSD PA_BBND PA_BNBD PA_NZ	BNSD BSND	(B, Q_N, Q_S)
BSND	BSND PA_BBND PA_BNBD PA_NZ	BSND	(B, Q_N, Q_S)
TND	TND PA_BBND PA_BNBD PA_NZ	TND	(Q_N, Q_T)

metadata validation

Parameter	Single-Parameter Validation	Existence Validation	Consistency Validation	Feature Cross-Validation
metadata	tensor_type only supports INT32. The shape is dynamically calculated by flash_attn_metadata. Not passing it is currently not supported. If it is not passed, an interception alarm is raised.	Optional Parameter	None	When passed, it must be consistent with the result generated by flash_attn_metadata.

Mask Parameter Group

mask_mode parameter description

• mask_mode=0: full computation mode (default value)
• mask_mode=3: Causal mode
• mask_mode=4: Window mode

Parameter	Single-Parameter Validation	Existence Validation	Consistency Validation	Feature Cross-Validation
mask_mode	data_type supports INT. Only supports input values 0, 3, and 4. The default value is 0.	Optional input. If this parameter is not passed, the default value is 0.	When mask_mode is 0, win_left and win_right must be -1, and attn_mask is not supported. When mask_mode is 3, win_left and win_right must be -1, and the attn_mask matrix must be passed. When mask_mode is 4, the attn_mask matrix must be passed.	None
attn_mask	tensor_type supports INT8. The tensor_shape dimension is (2048, 2048). This matrix is fixed as a 2048*2048 lower triangular matrix with all zeros on the diagonal.	None
win_left	data_type supports INT. If not passed or passed as -1, the value is positive infinity. The value must be >= -1.
win_right

SeqLengths Parameter Group

Parameter	Single-Parameter Validation	Existence Validation	Consistency Validation	Feature Cross-Validation
seqused_q	tensor_type only supports INT32. tensor_shape is (B,). Only supports non-negative integers. The value in seqused_q must be less than or equal to Q_S. The value in seqused_kv must be less than or equal to KV_S.	Optional Parameter	None	None
seqused_kv
cu_seqlens_q	tensor_type supports INT32. tensor_shape is (B+1,). Values support only non-negative integers. Values must be arranged in non-decreasing order (greater than or equal to the previous value). cu_seqlens_q: the first element is 0 and the last element equals Q_T. cu_seqlens_kv: the first element is 0 and the last element equals KV_T.	When layout_q is TND, cu_seqlens_q must be passed. When layout_q is not TND, cu_seqlens_q is not supported. When layout_kv is TND, cu_seqlens_kv must be passed. When layout_kv is not TND, cu_seqlens_kv is not supported.
cu_seqlens_kv
max_seqlen_q	data type supports INT. Default Value: -1.	Optional Parameter.	The value must be greater than or equal to -1. When passed, it must equal the actual maximum sequence length; otherwise, the behavior is undefined.
max_seqlen_kv

Paged Attention Parameter Group

Paged Attention is enabled when block_table is not empty.

Parameter	Single-Parameter Validation	Existence Validation	Consistency Validation	Feature Cross-Validation
block_table	tensor_type only supports INT32. tensor_shape is (B, max_num_blocks_per_seq). Values must be positive integers. 1024 >= block_size >= 16, and block_size % 16 == 0.	Optional Parameter	None	When PagedAttention is enabled, seqused_kv must be passed. When PagedAttention is enabled, block_table must not be empty.

Sinks Parameter Group

Parameter	Single-Parameter Validation	Existence Validation	Consistency Validation	Feature Cross-Validation
sinks	Not Supported.	Optional Parameter	None	Not supported yet.

SoftmaxLSE Parameter Group

Parameter	Single-Parameter Validation	Existence Validation	Consistency Validation	Feature Cross-Validation
return_softmax_lse	data_type only supports BOOL. The value only supports True and False. True enables softmax_lse, and False disables softmax_lse.	Optional attribute. Default Value: False.	When return_softmax_lse is False, outputs a Tensor with shape [1] and value 0. When return_softmax_lse is True, the output shape follows the layout matching relationship table.	None
softmax_lse	data_type only supports FLOAT32.	None		None

Example

• flash_attn_metadata + flash_attn joint call example (BSND)

  import torch
  import torch_npu
  import cann_ops_transformer
  
  torch_npu.npu.set_device(0)
  
  dtype = torch.bfloat16
  B = 2
  S = 16
  Q_N = 8
  KV_N = 2
  D = 128
  
  q = torch.randn(B, S, Q_N, D, dtype=dtype, device="npu")
  k = torch.randn(B, S, KV_N, D, dtype=dtype, device="npu")
  v = torch.randn(B, S, KV_N, D, dtype=dtype, device="npu")
  
  metadata = cann_ops_transformer.ops.flash_attn_metadata(
      Q_N,
      KV_N,
      D,
      batch_size=B,
      max_seqlen_q=S,
      max_seqlen_kv=S,
      mask_mode=0,
      win_left=-1,
      win_right=-1,
      layout_q="BSND",
      layout_kv="BSND",
      layout_out="BSND",
  )
  
  attn_out, softmax_lse = cann_ops_transformer.ops.flash_attn(
      q, k, v,
      metadata=metadata,
      softmax_scale=1.0 / (D ** 0.5),
      mask_mode=0,
      win_left=-1,
      win_right=-1,
      layout_q="BSND",
      layout_kv="BSND",
      layout_out="BSND",
      return_softmax_lse=False,
  )
  torch_npu.npu.synchronize()
  assert attn_out.shape == q.shape
  assert attn_out.dtype == q.dtype
  assert torch.isfinite(attn_out.float()).all().item()