lite_boost.layers.flash_attention

View Source On AtomGit
lite_boost.layers.flash_attention(q, k, v, q_lens=None, k_lens=None, dropout_p=0., softmax_scale=None, q_scale=None, causal=False, window_size=(-1, -1), deterministic=False, dtype=None, version=None)[source]

Computes flash attention with automatic backend fallback.

Selects FA3, FA2, or the NPU backend (npu_prompt_flash_attention) according to availability. Supports varlen sequences through q_lens/k_lens; when both are None every sequence is treated as full length.

Notation: B is the batch size, N the number of heads, D the head dim (D <= 256 on the NPU), S the sequence length of the input q/k/v tensors, lq the query sequence length (q.size(1), the same value as S), and lk the key sequence length (k.size(1)). The output always keeps the padded (B, S, N, D) layout of the input q (same shape as q) on every backend. In varlen mode q_lens/k_lens carry the per-sequence true lengths, which may be shorter than the full length and, on the NPU backend, may differ from each other across the batch; only the first q_lens[i] rows of sequence i are valid then – the rows beyond are masked kernel output with undefined values. When both are None every sequence is treated as full length (lq/lk).

Supported only on A2; 300I Duo is not supported. The dropout_p, causal, window_size, deterministic and version arguments are GPU-only features of the flash_attn (FA2/FA3) backends; the NPU backend ignores them and always computes global attention with no dropout.

Varlen support differs by backend: the GPU flash_attn (FA2/FA3) backends require every per-sequence length to equal the full length (q_lens[i] == lq and k_lens[i] == lk); the NPU backend additionally accepts shorter lengths, equal or unequal across the batch.

Parameters:

  • q (Tensor) – Query tensor with shape \((B, S, N, D)\), on an NPU device with head dim \(D \le 256\). Supported dtypes are float16, float32 and bfloat16.

  • k (Tensor) – Key tensor with shape \((B, S, N, D)\), same dtype as q.

  • v (Tensor) – Value tensor with shape \((B, S, N, D)\), same dtype as q.

  • q_lens (Union[list[int], Tensor[int32]], optional) – Per-sequence query lengths (length B) in varlen mode; may be shorter than the full length. Default: None.

  • k_lens (Union[list[int], Tensor[int32]], optional) – Per-sequence key lengths (length B) in varlen mode; may be shorter than the full key length. Default: None.

  • dropout_p (float, optional) – Dropout probability; effective on the GPU flash_attn FA2 backend only. Default: 0..

  • softmax_scale (float, optional) – Attention scaling factor; None means \(1/\sqrt{D}\). Default: None.

  • q_scale (float, optional) – Pre-scale applied to q as q = q * q_scale. Default: None.

  • causal (bool, optional) – Causal mask; effective on the GPU flash_attn FA2 backend only (the NPU backend always computes global attention, so True is not honored there). Default: False.

  • window_size (Tuple[int, int], optional) – Sliding-window restriction; effective on the GPU flash_attn FA2 backend only. Default: (-1, -1).

  • deterministic (bool, optional) – Deterministic mode; effective on the GPU flash_attn FA3/FA2 backends only. Default: False.

  • dtype (torch.dtype, optional) – Compute and output dtype, one of float16, bfloat16 or float32: q/k/v are cast to it before computing and the output is returned in it. None keeps the input dtype end to end (float16 in, float16 out); float32 inputs then still compute on the fp16 NPU kernel and come back as float32. Default: None.

  • version (int, optional) – 3 forces FA3, falling back to FA2 with a warning when FA3 is unavailable; None selects automatically. Effective on the GPU flash_attn FA3/FA2 backends only. Default: None.

Returns:

Tensor with the dtype requested via dtype (None keeps the input q dtype) and the same padded shape as the input q (\((B, S, N, D)\)) on every backend. In varlen mode only the first q_lens[i] rows of sequence i are valid; the rows beyond are masked kernel output with undefined values and should be sliced away (out[i, :q_lens[i]]).

Raises:

  • ValueError – If dtype is neither None nor in {float16, bfloat16, float32}, if q/k/v do not share a dtype in {float16, bfloat16, float32}, if q is not on an NPU device, if the head dim exceeds 256, if any per-sequence length exceeds the full length, or if q_lens/k_lens do not have one element per batch item (length B).

  • NotImplementedError – On the GPU flash_attn (FA2/FA3) backends, if varlen mode is used with lengths that do not equal the full length.

  • RuntimeError – If no backend is available (neither flash_attn nor NPU fusion attention).

Supported Platforms:

Ascend

Examples

>>> import torch
>>> import torch_npu
>>> from lite_boost.layers import flash_attention
>>> torch.npu.set_device(0)
>>> q = torch.randn(1, 16, 8, 32, device="npu")
>>> k = torch.randn(1, 16, 8, 32, device="npu")
>>> v = torch.randn(1, 16, 8, 32, device="npu")
>>> out = flash_attention(q, k, v)
>>> print(out.shape)
torch.Size([1, 16, 8, 32])
>>> print(out.dtype)
torch.float32