For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python function

flash_attention_ragged_gpu

flash_attention_ragged_gpu()

max.experimental.nn.common_layers.functional_kernels.flash_attention_ragged_gpu(q, k, v, input_row_offsets, max_seq_len, mask_variant, scale, local_window_size=-1)

source

Computes flash attention for ragged inputs using a GPU-optimized kernel, without a KV cache.

Parameters:

• q (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The query tensor, of shape [total_seq_len, num_heads, head_dim] (ragged).
• k (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The key tensor, of shape [total_seq_len, num_heads, head_dim] (ragged).
• v (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The value tensor, of shape [total_seq_len, num_heads, head_dim] (ragged).
• input_row_offsets (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The buffer of shape [batch_size + 1] with dtype uint32. Indicates where each sequence starts and ends in the ragged tensors. The values should be a prefix sum (cumulative sum) of sequence lengths.
• max_seq_len (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The maximum sequence length across the batch, as a rank-1 uint32 tensor on CPU.
• mask_variant (MHAMaskVariant) – The mask variant to use for attention.
• scale (float) – The scaling factor for attention scores.
• local_window_size (int) – The local window size for sliding window attention.

Returns:

The output tensor, of shape [total_seq_len, num_heads, head_dim].

Return type:

TensorValue