Struct MultiQueryAttention 

pub struct MultiQueryAttention {
    pub dim: usize,
    pub num_heads: usize,
    pub head_dim: usize,
    pub causal: bool,
    /* private fields */
}

Multi-Query Attention (MQA) - Extreme efficiency with single KV head

MQA uses a single key-value head shared across all query heads. This provides maximum memory efficiency at some quality cost.

Use Cases

• Inference-optimized models
• Memory-constrained deployments
• High-throughput serving

Example

let ctx = pollster::block_on(WgpuContext::new());
let engine = ComputeEngine::new(&ctx);
// [batch, seq_len, hidden]
let q = UnifiedTensor::zeros(&ctx, &[1, 128, 768]);
let k = UnifiedTensor::zeros(&ctx, &[1, 128, 64]);
let v = UnifiedTensor::zeros(&ctx, &[1, 128, 64]);
// 12 query heads share one KV head
let mqa = MultiQueryAttention::new(768, 12, true, 0.0);
let output = mqa.forward(&q, &k, &v, None, &ctx, &engine);

Fields

dim: usize

Total hidden dimension

num_heads: usize

Number of query heads

head_dim: usize

Dimension per head

causal: bool

Whether to use causal masking

Implementations

impl MultiQueryAttention

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

Create a new MQA layer.

Arguments

• dim - Total hidden dimension
• num_heads - Number of query heads
• causal - Whether to apply causal masking

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

Forward pass with multi-query attention.

Arguments

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

Returns

Output tensor [batch, seq_q, num_heads, head_dim]

pub fn memory_savings(&self) -> String

Get memory savings compared to MHA

Trait Implementations

impl Debug for MultiQueryAttention

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

Formats the value using the given formatter.