Class GroupedQueryAttention

java.lang.Object
smile.llm.model.llama.GroupedQueryAttention
All Implemented Interfaces:
Attention
public class GroupedQueryAttention extends Object implements Attention
Grouped Query Attention (GQA). GQA is a highly efficient transformer attention mechanism that bridges the gap between traditional Multi-Head Attention (MHA) and Multi-Query Attention (MQA). By grouping sets of query heads to share a single Key and Value head, GQA drastically reduces memory usage during inference while maintaining model quality.

KV activations are stored in a shared KvCachePool managed by the inference engine (e.g. smile-serve), enabling radix-tree prefix reuse across requests instead of allocating a private maxBatchSize × maxSeqLen buffer per layer.

  • Constructor Details

    • GroupedQueryAttention

      public GroupedQueryAttention(int dim, int numHeads, int numKvHeads, int layerId)
      Constructor. Install a KvCachePool via setCachePool(KvCachePool) before the first forward(Tensor, int, Tensor, Tensor).
      Parameters:
      dim - token embedding dimension.
      numHeads - number of query heads.
      numKvHeads - number of key/value heads.
      layerId - zero-based layer index within the transformer.

    • GroupedQueryAttention

      public GroupedQueryAttention(int dim, int numHeads, int numKvHeads, KvCacheLayout layout)
      Convenience constructor that allocates a small test pool sized from layout. Prefer installing a shared pool in production.
      Parameters:
      dim - token embedding dimension.
      numHeads - number of query heads.
      numKvHeads - number of key/value heads.
      layout - cache layout used for the private test pool.

  • Method Details

    • replaceProjections

      public void replaceProjections(LinearOp wq, LinearOp wk, LinearOp wv, LinearOp wo)
      Replaces projection ops with quantized backends (shard-then-pack already applied). Unregisters and frees dense LinearLayer shells so their GPU storage is released before KV pool sizing.
      Parameters:
      wq - query projection op.
      wk - key projection op.
      wv - value projection op.
      wo - output projection op.

    • numQueryHeads

      public int numQueryHeads()
      Returns the query head count on this rank.
      Returns:
      query head count (equals global heads when TP size is 1).

    • numKeyValueHeads

      public int numKeyValueHeads()
      Returns the key/value head count.
      Returns:
      key/value head count (global; local when TP is applied).

    • module

      public MemorySegment module()
      Description copied from interface: Attention
      Returns pytorch module.
      Specified by:
      module in interface Attention
      Returns:
      pytorch module.

    • forward

      public Tensor forward(Tensor x, int startPos, Tensor cis, Tensor mask)
      Description copied from interface: Attention
      Forward pass through the attention module.
      Specified by:
      forward in interface Attention
      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

      public Tensor forward(Tensor x, int[] positions, Tensor cis, Tensor mask)
      Description copied from interface: Attention
      Decode forward with per-row cache write positions (FlashInfer ragged path).

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

      Specified by:
      forward in interface Attention
      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.