% Generated by roxygen2: do not edit by hand
% Please edit documentation in R/operations.R
\name{ggml_flash_attn_ext}
\alias{ggml_flash_attn_ext}
\title{Flash Attention (Graph)}
\usage{
ggml_flash_attn_ext(
  ctx,
  q,
  k,
  v,
  mask = NULL,
  scale,
  max_bias = 0,
  logit_softcap = 0
)
}
\arguments{
\item{ctx}{GGML context}

\item{q}{Query tensor of shape [head_dim, n_head, n_tokens, batch]}

\item{k}{Key tensor of shape [head_dim, n_head_kv, n_kv, batch]}

\item{v}{Value tensor of shape [head_dim, n_head_kv, n_kv, batch]}

\item{mask}{Optional attention mask tensor (NULL for no mask).
For causal attention, use ggml_diag_mask_inf instead.}

\item{scale}{Attention scale factor, typically 1/sqrt(head_dim)}

\item{max_bias}{Maximum ALiBi bias (0.0 to disable ALiBi)}

\item{logit_softcap}{Logit soft-capping value (0.0 to disable).
Used by some models like Gemma 2.}
}
\value{
Attention output tensor of shape [head_dim, n_head, n_tokens, batch]
}
\description{
Creates a graph node for Flash Attention computation.
This is a memory-efficient implementation of scaled dot-product attention.
}
\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{
\donttest{
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)
}
}
