Interface Attention

All Known Implementing Classes:
GatedAttention, GroupedQueryAttention
public interface Attention
Multi-head attention. Multi-head attention is a core component of the Transformer deep learning architecture. It allows AI models to analyze data simultaneously from multiple representation subspaces, improving their ability to capture complex, varied relationships between words or tokens in a sequence.

Instead of performing a single attention function, multi-head attention linearly projects queries, keys, and values into multiple smaller dimensions. These projections are processed in parallel by distinct attention heads.

apply(Tensor, Tensor, Tensor, Tensor) dispatches to the process-wide AttentionBackends kernel (flashinfer by default when available, else torch_native).

  • Method Summary

    Modifier and Type
    Method
    Description
    default Tensor
    apply(Tensor query, Tensor key, Tensor value, Tensor mask)
    Computes the scaled dot product attention on query, key and value tensors, using an optional attention mask if passed, and applying dropout if a probability greater than 0.0 is specified.
    default Tensor
    apply(Tensor query, Tensor key, Tensor value, Tensor mask, double dropout, boolean isCausal, double scale)
    Computes the scaled dot product attention on query, key and value tensors, using an optional attention mask if passed, and applying dropout if a probability greater than 0.0 is specified.
    default Tensor
    forward(Tensor x, int[] positions, Tensor cis, Tensor mask)
    Decode forward with per-row cache write positions (FlashInfer ragged path).
    Tensor
    forward(Tensor x, int startPos, Tensor cis, Tensor mask)
    Forward pass through the attention module.
    MemorySegment
    module()
    Returns pytorch module.
    default Tensor
    repeatKV(Tensor input, int numRep)
    Efficiently repeat a tensor.

  • Method Details

    • forward

      Tensor forward(Tensor x, int startPos, Tensor cis, Tensor mask)
      Forward pass through the attention module.
      Parameters:
      x - the input tensor.
      startPos - the starting position for attention caching.
      cis - the precomputed frequency tensor.
      mask - the attention mask tensor.
      Returns:
      the output tensor.

    • forward

      default Tensor forward(Tensor x, int[] positions, Tensor cis, Tensor mask)
      Decode forward with per-row cache write positions (FlashInfer ragged path).

      Default implementation rejects mixed positions; backends that support ragged decode override this.

      Parameters:
      x - input [B, 1, D].
      positions - absolute write position per batch row.
      cis - RoPE freqs [B, headDim/2] (gathered) or uniform.
      mask - attention mask, or null for decode.
      Returns:
      attention output.

    • module

      MemorySegment module()
      Returns pytorch module.
      Returns:
      pytorch module.

    • apply

      default Tensor apply(Tensor query, Tensor key, Tensor value, Tensor mask)
      Computes the scaled dot product attention on query, key and value tensors, using an optional attention mask if passed, and applying dropout if a probability greater than 0.0 is specified.
      Parameters:
      query - the query tensor.
      key - the key tensor.
      value - the value tensor.
      mask - the attention mask.
      Returns:
      the attention output.

    • apply

      default Tensor apply(Tensor query, Tensor key, Tensor value, Tensor mask, double dropout, boolean isCausal, double scale)
      Computes the scaled dot product attention on query, key and value tensors, using an optional attention mask if passed, and applying dropout if a probability greater than 0.0 is specified.
      Parameters:
      query - the query tensor.
      key - the key tensor.
      value - the value tensor.
      mask - the attention mask.
      dropout - the dropout probability; if greater than 0.0, dropout is applied.
      isCausal - If set to true, the attention masking is a lower triangular matrix when the mask is a square matrix. The attention masking has the form of the upper left causal bias due to the alignment when the mask is a non-square matrix. An error is thrown if both attn_mask and is_causal are set.
      scale - Optional scaling factor applied prior to softmax. If ≤ 0, the standard scaling factor is used.
      Returns:
      the attention output.

    • repeatKV

      default Tensor repeatKV(Tensor input, int numRep)
      Efficiently repeat a tensor.
      Parameters:
      input - the input tensor to repeat.
      numRep - the number of times to repeat.
      Returns:
      the repeated tensor.