# Flash Attention (Graph) Creates a graph node for Flash Attention computation. This is a memory-efficient implementation of scaled dot-product attention. ```r ggml_flash_attn_ext( ctx, q, k, v, mask = NULL, scale, max_bias = 0, logit_softcap = 0 ) ``` ## Arguments - `ctx`: GGML context - `q`: Query tensor of shape [head_dim, n_head, n_tokens, batch] - `k`: Key tensor of shape [head_dim, n_head_kv, n_kv, batch] - `v`: Value tensor of shape [head_dim, n_head_kv, n_kv, batch] - `mask`: Optional attention mask tensor (NULL for no mask). For causal attention, use ggml_diag_mask_inf instead. - `scale`: Attention scale factor, typically 1/sqrt(head_dim) - `max_bias`: Maximum ALiBi bias (0.0 to disable ALiBi) - `logit_softcap`: Logit soft-capping value (0.0 to disable). Used by some models like Gemma 2. ## Returns Attention output tensor of shape [head_dim, n_head, n_tokens, batch] ## Details Flash Attention computes: softmax(Q * K^T / scale + mask) * V Key features: - Memory efficient: O(n) instead of O(n^2) memory for attention matrix - Supports grouped-query attention (GQA) when n_head_kv < n_head - Supports multi-query attention (MQA) when n_head_kv = 1 - Optional ALiBi (Attention with Linear Biases) for position encoding - Optional logit soft-capping for numerical stability ## Examples ```r ctx <- ggml_init(64 * 1024 * 1024) head_dim <- 64 n_head <- 8 n_head_kv <- 2 # GQA with 4:1 ratio seq_len <- 32 q <- ggml_new_tensor_4d(ctx, GGML_TYPE_F32, head_dim, n_head, seq_len, 1) k <- ggml_new_tensor_4d(ctx, GGML_TYPE_F32, head_dim, n_head_kv, seq_len, 1) v <- ggml_new_tensor_4d(ctx, GGML_TYPE_F32, head_dim, n_head_kv, seq_len, 1) ggml_set_f32(q, rnorm(head_dim * n_head * seq_len)) ggml_set_f32(k, rnorm(head_dim * n_head_kv * seq_len)) ggml_set_f32(v, rnorm(head_dim * n_head_kv * seq_len)) # Scale = 1/sqrt(head_dim) scale <- 1.0 / sqrt(head_dim) # Compute attention out <- ggml_flash_attn_ext(ctx, q, k, v, NULL, scale, 0.0, 0.0) graph <- ggml_build_forward_expand(ctx, out) ggml_graph_compute(ctx, graph) ggml_free(ctx) ```

ggml_flash_attn_ext function