Struct GroupedQueryAttention 

pub struct GroupedQueryAttention {
    pub dim: usize,
    pub num_heads: usize,
    pub num_kv_heads: usize,
    pub head_dim: usize,
    pub num_groups: usize,
    pub causal: bool,
    pub dropout_prob: f32,
    /* private fields */
}

Grouped Query Attention (GQA) - Efficient attention with fewer KV heads

GQA uses fewer key-value heads than query heads, reducing memory and computation while maintaining quality. This is the attention mechanism used by Llama 3, Mistral, and many modern LLMs.

Memory Savings

• Standard MHA: num_heads K and V projections
• GQA: num_kv_heads K and V projections (typically 2-8x fewer)

Example

let ctx = pollster::block_on(WgpuContext::new());
let engine = ComputeEngine::new(&ctx);
// [batch, seq_len, hidden]
let q = UnifiedTensor::zeros(&ctx, &[1, 128, 4096]);
let k = UnifiedTensor::zeros(&ctx, &[1, 128, 1024]);
let v = UnifiedTensor::zeros(&ctx, &[1, 128, 1024]);
// Llama 3 8B: 32 query heads share 8 KV heads (no dropout, no mask)
let gqa = GroupedQueryAttention::new(4096, 32, 8, true, 0.0);
let output = gqa.forward(&q, &k, &v, None, &ctx, &engine);

Fields

dim: usize

Total hidden dimension

num_heads: usize

Number of query heads

num_kv_heads: usize

Number of key-value heads (fewer than query heads)

head_dim: usize

Dimension per head

num_groups: usize

Number of query heads per KV head

causal: bool

Whether to use causal masking

dropout_prob: f32

Dropout probability (0.0 = no dropout)

Implementations

impl GroupedQueryAttention

pub fn new( dim: usize, num_heads: usize, num_kv_heads: usize, causal: bool, dropout_prob: f32, ) -> Self

Create a new GQA layer.

Arguments

• dim - Total hidden dimension
• num_heads - Number of query heads
• num_kv_heads - Number of key-value heads (must divide num_heads evenly)
• causal - Whether to apply causal masking
• dropout_prob - Dropout probability

Panics

Panics if num_heads is not divisible by num_kv_heads

pub fn forward( &self, q: &UnifiedTensor, k: &UnifiedTensor, v: &UnifiedTensor, mask: Option<&UnifiedTensor>, ctx: &WgpuContext, engine: &ComputeEngine, ) -> UnifiedTensor

Forward pass with grouped query attention.

Arguments

• q - Query tensor [batch, seq_q, num_heads, head_dim]
• k - Key tensor [batch, seq_k, num_kv_heads, head_dim]
• v - Value tensor [batch, seq_k, num_kv_heads, head_dim]

Returns

Output tensor [batch, seq_q, num_heads, head_dim]

pub fn forward_at( &self, q: &UnifiedTensor, k: &UnifiedTensor, v: &UnifiedTensor, mask: Option<&UnifiedTensor>, pos: usize, ctx: &WgpuContext, engine: &ComputeEngine, ) -> UnifiedTensor

Like forward, with the first query at absolute position pos: query i attends to keys 0..=pos + i when causal. Used to decode against a KV cache, where k/v hold the cache.

pub fn memory_savings(&self) -> String

Get memory savings compared to standard MHA

Trait Implementations

impl Debug for GroupedQueryAttention

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.