osobhandClaude Sonnet 5
4aaa36a57a
Whole-workspace rustfmt pass picked up while iterating on Mamba GPU backward work. Verified formatting-only via diff sampling; no logic changed. Co-Authored-By: Claude Sonnet 5 <[email protected]>
661 lines
25 KiB
Rust
661 lines
25 KiB
Rust
//! Grouped Query Attention (GQA) KV head expansion for inference.
|
||
//!
|
||
//! Modern LLMs such as Llama 2/3, Mistral, Gemma, and Qwen use GQA where
|
||
//! `num_kv_heads < num_q_heads`. During the attention forward pass each KV
|
||
//! head is shared by `queries_per_group = num_q_heads / num_kv_heads` query
|
||
//! heads, eliminating the need to store a full per-head KV cache.
|
||
//!
|
||
//! This module provides:
|
||
//! - [`GqaConfig`] — validated configuration object.
|
||
//! - [`expand_kv_heads`] — repeat/tile KV tensor to match Q head count.
|
||
//! - [`kv_head_for_q`] — zero-cost index mapping from Q head → KV head.
|
||
//! - [`gqa_attention_cpu`] — single-batch GQA forward pass on the CPU.
|
||
//!
|
||
//! # Layout convention
|
||
//!
|
||
//! Tensors are represented as flat `&[f32]` / `Vec<f32>` in
|
||
//! **row-major (C-contiguous)** order. The logical shape is written in
|
||
//! square brackets throughout the docs.
|
||
//!
|
||
//! ```text
|
||
//! KV tensor: [batch, num_kv_heads, seq_len, head_dim]
|
||
//! Q tensor: [num_q_heads, seq_len_q, head_dim] (single batch item)
|
||
//! ```
|
||
|
||
use thiserror::Error;
|
||
|
||
// ──────────────────────────────────────────────────────────────────────────────
|
||
// Error type
|
||
// ──────────────────────────────────────────────────────────────────────────────
|
||
|
||
/// Errors that can occur when constructing or using a [`GqaConfig`].
|
||
#[derive(Debug, Error, PartialEq, Eq)]
|
||
pub enum GqaError {
|
||
/// `num_q_heads` is not evenly divisible by `num_kv_heads`.
|
||
#[error("num_q_heads ({q}) must be divisible by num_kv_heads ({kv})")]
|
||
IndivisibleHeads { q: usize, kv: usize },
|
||
|
||
/// Either head count is zero.
|
||
#[error("num_q_heads and num_kv_heads must be > 0")]
|
||
ZeroHeads,
|
||
|
||
/// `num_kv_heads` is larger than `num_q_heads`, which is not valid.
|
||
#[error("num_kv_heads ({kv}) cannot exceed num_q_heads ({q})")]
|
||
KvExceedsQ { q: usize, kv: usize },
|
||
|
||
/// `head_dim` is zero.
|
||
#[error("head_dim must be > 0")]
|
||
ZeroHeadDim,
|
||
}
|
||
|
||
// ──────────────────────────────────────────────────────────────────────────────
|
||
// Configuration
|
||
// ──────────────────────────────────────────────────────────────────────────────
|
||
|
||
/// Configuration for Grouped Query Attention.
|
||
///
|
||
/// # Invariants (upheld by [`GqaConfig::new`])
|
||
/// - `num_q_heads > 0`
|
||
/// - `num_kv_heads > 0`
|
||
/// - `num_kv_heads <= num_q_heads`
|
||
/// - `num_q_heads % num_kv_heads == 0`
|
||
/// - `head_dim > 0`
|
||
///
|
||
/// # Examples
|
||
///
|
||
/// ```
|
||
/// use rtx_inference::gqa::GqaConfig;
|
||
///
|
||
/// // Llama 3 8B: 32 Q heads, 8 KV heads, head_dim 128.
|
||
/// let cfg = GqaConfig::new(32, 8, 128).unwrap();
|
||
/// assert_eq!(cfg.queries_per_group(), 4);
|
||
/// assert!(cfg.is_grouped());
|
||
/// assert!(!cfg.is_mqa());
|
||
/// assert!((cfg.kv_memory_ratio() - 0.25).abs() < 1e-6);
|
||
/// ```
|
||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||
pub struct GqaConfig {
|
||
/// Total number of query heads.
|
||
pub num_q_heads: usize,
|
||
/// Total number of key/value heads (≤ `num_q_heads`).
|
||
pub num_kv_heads: usize,
|
||
/// Dimension of each head.
|
||
pub head_dim: usize,
|
||
}
|
||
|
||
impl GqaConfig {
|
||
/// Construct a validated [`GqaConfig`].
|
||
///
|
||
/// # Errors
|
||
///
|
||
/// Returns [`GqaError`] when any invariant is violated.
|
||
pub fn new(num_q_heads: usize, num_kv_heads: usize, head_dim: usize) -> Result<Self, GqaError> {
|
||
if num_q_heads == 0 || num_kv_heads == 0 {
|
||
return Err(GqaError::ZeroHeads);
|
||
}
|
||
if head_dim == 0 {
|
||
return Err(GqaError::ZeroHeadDim);
|
||
}
|
||
if num_kv_heads > num_q_heads {
|
||
return Err(GqaError::KvExceedsQ {
|
||
q: num_q_heads,
|
||
kv: num_kv_heads,
|
||
});
|
||
}
|
||
if num_q_heads % num_kv_heads != 0 {
|
||
return Err(GqaError::IndivisibleHeads {
|
||
q: num_q_heads,
|
||
kv: num_kv_heads,
|
||
});
|
||
}
|
||
Ok(Self {
|
||
num_q_heads,
|
||
num_kv_heads,
|
||
head_dim,
|
||
})
|
||
}
|
||
|
||
/// Number of query heads that share one KV head.
|
||
///
|
||
/// Equals `num_q_heads / num_kv_heads`. For standard MHA this is `1`.
|
||
#[inline]
|
||
pub fn queries_per_group(&self) -> usize {
|
||
self.num_q_heads / self.num_kv_heads
|
||
}
|
||
|
||
/// Returns `true` when GQA is active (`num_kv_heads < num_q_heads`).
|
||
#[inline]
|
||
pub fn is_grouped(&self) -> bool {
|
||
self.num_kv_heads < self.num_q_heads
|
||
}
|
||
|
||
/// Returns `true` for Multi-Query Attention (single KV head for all Q heads).
|
||
#[inline]
|
||
pub fn is_mqa(&self) -> bool {
|
||
self.num_kv_heads == 1
|
||
}
|
||
|
||
/// Ratio of KV memory relative to full MHA.
|
||
///
|
||
/// E.g. 8 KV heads / 32 Q heads = 0.25 → 75 % memory saving.
|
||
#[inline]
|
||
pub fn kv_memory_ratio(&self) -> f32 {
|
||
self.num_kv_heads as f32 / self.num_q_heads as f32
|
||
}
|
||
}
|
||
|
||
// ──────────────────────────────────────────────────────────────────────────────
|
||
// Index mapping
|
||
// ──────────────────────────────────────────────────────────────────────────────
|
||
|
||
/// Map a query head index to its corresponding KV head index.
|
||
///
|
||
/// The mapping is: `kv_head = q_head / queries_per_group`.
|
||
///
|
||
/// This is a zero-cost inline function suitable for tight inner loops.
|
||
///
|
||
/// # Examples
|
||
///
|
||
/// ```
|
||
/// use rtx_inference::gqa::kv_head_for_q;
|
||
///
|
||
/// // 8 Q heads, 2 KV heads → queries_per_group = 4
|
||
/// assert_eq!(kv_head_for_q(0, 4), 0);
|
||
/// assert_eq!(kv_head_for_q(3, 4), 0);
|
||
/// assert_eq!(kv_head_for_q(4, 4), 1);
|
||
/// assert_eq!(kv_head_for_q(7, 4), 1);
|
||
/// ```
|
||
#[inline]
|
||
pub fn kv_head_for_q(q_head: usize, queries_per_group: usize) -> usize {
|
||
q_head / queries_per_group
|
||
}
|
||
|
||
// ──────────────────────────────────────────────────────────────────────────────
|
||
// KV head expansion
|
||
// ──────────────────────────────────────────────────────────────────────────────
|
||
|
||
/// Expand KV heads to match the query head count by repetition (tiling).
|
||
///
|
||
/// Input layout: `[batch, num_kv_heads, seq_len, head_dim]`
|
||
/// Output layout: `[batch, num_q_heads, seq_len, head_dim]`
|
||
///
|
||
/// Each KV head `i` is repeated [`GqaConfig::queries_per_group`] times
|
||
/// consecutively so that the output can be used directly with a full-MHA
|
||
/// kernel.
|
||
///
|
||
/// # Memory
|
||
///
|
||
/// Allocates a new `Vec<f32>` of length
|
||
/// `batch_size * num_q_heads * seq_len * head_dim`.
|
||
///
|
||
/// # Panics
|
||
///
|
||
/// Panics in debug mode if `kv.len()` does not match the expected size
|
||
/// `batch_size * config.num_kv_heads * seq_len * config.head_dim`.
|
||
///
|
||
/// # Examples
|
||
///
|
||
/// ```
|
||
/// use rtx_inference::gqa::{GqaConfig, expand_kv_heads};
|
||
///
|
||
/// let cfg = GqaConfig::new(8, 2, 4).unwrap();
|
||
/// // KV: [1, 2, 3, 4] (batch=1, num_kv=2, seq=3, head_dim=4)
|
||
/// let kv: Vec<f32> = (0..24).map(|x| x as f32).collect();
|
||
/// let out = expand_kv_heads(&kv, &cfg, 1, 3);
|
||
/// assert_eq!(out.len(), 1 * 8 * 3 * 4);
|
||
/// ```
|
||
pub fn expand_kv_heads(
|
||
kv: &[f32],
|
||
config: &GqaConfig,
|
||
batch_size: usize,
|
||
seq_len: usize,
|
||
) -> Vec<f32> {
|
||
let kv_head_stride = seq_len * config.head_dim;
|
||
let kv_batch_stride = config.num_kv_heads * kv_head_stride;
|
||
let q_head_stride = seq_len * config.head_dim;
|
||
let q_batch_stride = config.num_q_heads * q_head_stride;
|
||
let gpg = config.queries_per_group();
|
||
|
||
debug_assert_eq!(
|
||
kv.len(),
|
||
batch_size * kv_batch_stride,
|
||
"KV slice length mismatch: expected {} got {}",
|
||
batch_size * kv_batch_stride,
|
||
kv.len()
|
||
);
|
||
|
||
let mut out = vec![0.0f32; batch_size * q_batch_stride];
|
||
|
||
for b in 0..batch_size {
|
||
for q_head in 0..config.num_q_heads {
|
||
let kv_head = kv_head_for_q(q_head, gpg);
|
||
let kv_offset = b * kv_batch_stride + kv_head * kv_head_stride;
|
||
let out_offset = b * q_batch_stride + q_head * q_head_stride;
|
||
// Copy the entire [seq_len, head_dim] slice.
|
||
out[out_offset..out_offset + kv_head_stride]
|
||
.copy_from_slice(&kv[kv_offset..kv_offset + kv_head_stride]);
|
||
}
|
||
}
|
||
|
||
out
|
||
}
|
||
|
||
// ──────────────────────────────────────────────────────────────────────────────
|
||
// CPU attention forward pass
|
||
// ──────────────────────────────────────────────────────────────────────────────
|
||
|
||
/// Compute GQA scaled dot-product attention for a single batch item on the CPU.
|
||
///
|
||
/// Tensors are flat row-major slices:
|
||
/// - `q`: `[num_q_heads, seq_len_q, head_dim]`
|
||
/// - `k`: `[num_kv_heads, seq_len_kv, head_dim]` *(NOT pre-expanded)*
|
||
/// - `v`: `[num_kv_heads, seq_len_kv, head_dim]` *(NOT pre-expanded)*
|
||
///
|
||
/// Returns `[num_q_heads, seq_len_q, head_dim]`.
|
||
///
|
||
/// For each query head `h_q`, the corresponding KV head is
|
||
/// `h_q / queries_per_group`. Attention is computed as:
|
||
///
|
||
/// ```text
|
||
/// scores[h_q, i, j] = dot(Q[h_q, i, :], K[h_kv, j, :]) * scale
|
||
/// attn[h_q, i, j] = softmax(scores[h_q, i, :])[j] (causal mask applied before softmax)
|
||
/// out[h_q, i, :] = sum_j attn[h_q, i, j] * V[h_kv, j, :]
|
||
/// ```
|
||
///
|
||
/// When `causal = true` positions `j > i` are masked to `-∞` before softmax.
|
||
///
|
||
/// # Panics
|
||
///
|
||
/// Panics in debug mode on slice length mismatch.
|
||
///
|
||
/// # Examples
|
||
///
|
||
/// ```
|
||
/// use rtx_inference::gqa::{GqaConfig, gqa_attention_cpu};
|
||
///
|
||
/// let cfg = GqaConfig::new(2, 2, 4).unwrap();
|
||
/// let scale = (4f32).sqrt().recip();
|
||
/// let q: Vec<f32> = (0..16).map(|x| x as f32 * 0.01).collect();
|
||
/// let k = q.clone();
|
||
/// let v = q.clone();
|
||
/// let out = gqa_attention_cpu(&q, &k, &v, &cfg, 2, 2, scale, false);
|
||
/// assert_eq!(out.len(), 2 * 2 * 4);
|
||
/// ```
|
||
pub fn gqa_attention_cpu(
|
||
q: &[f32],
|
||
k: &[f32],
|
||
v: &[f32],
|
||
config: &GqaConfig,
|
||
seq_len_q: usize,
|
||
seq_len_kv: usize,
|
||
scale: f32,
|
||
causal: bool,
|
||
) -> Vec<f32> {
|
||
let head_dim = config.head_dim;
|
||
let gpg = config.queries_per_group();
|
||
|
||
debug_assert_eq!(q.len(), config.num_q_heads * seq_len_q * head_dim);
|
||
debug_assert_eq!(k.len(), config.num_kv_heads * seq_len_kv * head_dim);
|
||
debug_assert_eq!(v.len(), config.num_kv_heads * seq_len_kv * head_dim);
|
||
|
||
// Strides for input tensors.
|
||
let q_head_stride = seq_len_q * head_dim;
|
||
let kv_head_stride = seq_len_kv * head_dim;
|
||
|
||
// Output: [num_q_heads, seq_len_q, head_dim]
|
||
let mut output = vec![0.0f32; config.num_q_heads * seq_len_q * head_dim];
|
||
|
||
// Scratch buffer for attention scores of one query position: [seq_len_kv]
|
||
let mut scores = vec![0.0f32; seq_len_kv];
|
||
|
||
for h_q in 0..config.num_q_heads {
|
||
let h_kv = kv_head_for_q(h_q, gpg);
|
||
let q_base = h_q * q_head_stride;
|
||
let kv_base = h_kv * kv_head_stride;
|
||
let out_base = h_q * q_head_stride;
|
||
|
||
for i in 0..seq_len_q {
|
||
// Compute raw dot-product scores for query position i.
|
||
let q_row = &q[q_base + i * head_dim..q_base + i * head_dim + head_dim];
|
||
|
||
let max_j = if causal { i + 1 } else { seq_len_kv };
|
||
|
||
// Fill scores; positions beyond causal boundary stay at -∞.
|
||
for j in 0..seq_len_kv {
|
||
if j < max_j {
|
||
let k_row = &k[kv_base + j * head_dim..kv_base + j * head_dim + head_dim];
|
||
scores[j] = dot(q_row, k_row) * scale;
|
||
} else {
|
||
scores[j] = f32::NEG_INFINITY;
|
||
}
|
||
}
|
||
|
||
// Numerically stable softmax over scores[0..seq_len_kv].
|
||
let max_score = scores[..seq_len_kv]
|
||
.iter()
|
||
.copied()
|
||
.fold(f32::NEG_INFINITY, f32::max);
|
||
|
||
let mut sum_exp = 0.0f32;
|
||
for s in &mut scores[..seq_len_kv] {
|
||
*s = (*s - max_score).exp();
|
||
sum_exp += *s;
|
||
}
|
||
// Guard against the degenerate all-masked case.
|
||
let inv_sum = if sum_exp > 0.0 { sum_exp.recip() } else { 0.0 };
|
||
for s in &mut scores[..seq_len_kv] {
|
||
*s *= inv_sum;
|
||
}
|
||
|
||
// Weighted sum over value vectors.
|
||
let out_row = &mut output[out_base + i * head_dim..out_base + i * head_dim + head_dim];
|
||
for j in 0..seq_len_kv {
|
||
let attn = scores[j];
|
||
if attn == 0.0 {
|
||
continue;
|
||
}
|
||
let v_row = &v[kv_base + j * head_dim..kv_base + j * head_dim + head_dim];
|
||
for d in 0..head_dim {
|
||
out_row[d] += attn * v_row[d];
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
output
|
||
}
|
||
|
||
// ──────────────────────────────────────────────────────────────────────────────
|
||
// Internal helpers
|
||
// ──────────────────────────────────────────────────────────────────────────────
|
||
|
||
/// Dot product of two equal-length slices.
|
||
#[inline]
|
||
fn dot(a: &[f32], b: &[f32]) -> f32 {
|
||
debug_assert_eq!(a.len(), b.len());
|
||
a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()
|
||
}
|
||
|
||
// ──────────────────────────────────────────────────────────────────────────────
|
||
// Tests
|
||
// ──────────────────────────────────────────────────────────────────────────────
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
// ── Config validation tests ───────────────────────────────────────────────
|
||
|
||
/// MHA (num_q == num_kv) → queries_per_group == 1.
|
||
#[test]
|
||
fn test_mha_config_queries_per_group_one() {
|
||
let cfg = GqaConfig::new(8, 8, 64).unwrap();
|
||
assert_eq!(cfg.queries_per_group(), 1);
|
||
assert!(!cfg.is_grouped());
|
||
}
|
||
|
||
/// num_q not divisible by num_kv → IndivisibleHeads error.
|
||
#[test]
|
||
fn test_gqa_config_indivisible_error() {
|
||
let err = GqaConfig::new(7, 3, 64).unwrap_err();
|
||
assert_eq!(err, GqaError::IndivisibleHeads { q: 7, kv: 3 });
|
||
}
|
||
|
||
/// num_kv > num_q → KvExceedsQ error.
|
||
#[test]
|
||
fn test_gqa_config_kv_exceeds_q_error() {
|
||
let err = GqaConfig::new(4, 8, 64).unwrap_err();
|
||
assert_eq!(err, GqaError::KvExceedsQ { q: 4, kv: 8 });
|
||
}
|
||
|
||
/// num_q == 0 → ZeroHeads error.
|
||
#[test]
|
||
fn test_gqa_config_zero_heads_error() {
|
||
let err = GqaConfig::new(0, 0, 64).unwrap_err();
|
||
assert_eq!(err, GqaError::ZeroHeads);
|
||
}
|
||
|
||
/// num_kv == 0 with non-zero num_q → ZeroHeads error.
|
||
#[test]
|
||
fn test_gqa_config_zero_kv_heads_error() {
|
||
let err = GqaConfig::new(8, 0, 64).unwrap_err();
|
||
assert_eq!(err, GqaError::ZeroHeads);
|
||
}
|
||
|
||
/// num_kv < num_q → is_grouped() == true.
|
||
#[test]
|
||
fn test_is_grouped_true() {
|
||
let cfg = GqaConfig::new(32, 8, 128).unwrap();
|
||
assert!(cfg.is_grouped());
|
||
}
|
||
|
||
/// num_kv == 1 → is_mqa() == true.
|
||
#[test]
|
||
fn test_is_mqa() {
|
||
let cfg = GqaConfig::new(8, 1, 64).unwrap();
|
||
assert!(cfg.is_mqa());
|
||
assert!(cfg.is_grouped());
|
||
}
|
||
|
||
/// 8 kv / 32 q → kv_memory_ratio == 0.25.
|
||
#[test]
|
||
fn test_kv_memory_ratio() {
|
||
let cfg = GqaConfig::new(32, 8, 128).unwrap();
|
||
assert!((cfg.kv_memory_ratio() - 0.25_f32).abs() < 1e-6);
|
||
}
|
||
|
||
// ── Index mapping tests ───────────────────────────────────────────────────
|
||
|
||
/// q_head=5, queries_per_group=4 → kv_head=1.
|
||
#[test]
|
||
fn test_kv_head_for_q_mapping() {
|
||
assert_eq!(kv_head_for_q(5, 4), 1);
|
||
}
|
||
|
||
/// Full group boundary checks for queries_per_group=4.
|
||
#[test]
|
||
fn test_kv_head_for_q_boundary() {
|
||
// 8 Q heads, 2 KV heads, gpg=4
|
||
for q in 0..4usize {
|
||
assert_eq!(kv_head_for_q(q, 4), 0);
|
||
}
|
||
for q in 4..8usize {
|
||
assert_eq!(kv_head_for_q(q, 4), 1);
|
||
}
|
||
}
|
||
|
||
// ── expand_kv_heads shape test ────────────────────────────────────────────
|
||
|
||
/// Input [1, 2, 4, 8] → output [1, 8, 4, 8] for num_q=8, num_kv=2.
|
||
#[test]
|
||
fn test_expand_kv_heads_shape() {
|
||
let cfg = GqaConfig::new(8, 2, 8).unwrap();
|
||
let batch = 1;
|
||
let seq = 4;
|
||
let kv: Vec<f32> = (0..(batch * 2 * seq * 8)).map(|x| x as f32).collect();
|
||
let out = expand_kv_heads(&kv, &cfg, batch, seq);
|
||
assert_eq!(out.len(), batch * 8 * seq * 8);
|
||
}
|
||
|
||
/// KV[0] is repeated 4× (first four Q slots), KV[1] repeats for the rest.
|
||
#[test]
|
||
fn test_expand_kv_heads_values() {
|
||
// cfg: 8 Q heads, 2 KV heads, head_dim=4, batch=1, seq=1
|
||
let cfg = GqaConfig::new(8, 2, 4).unwrap();
|
||
// KV head 0: [1,2,3,4], KV head 1: [5,6,7,8]
|
||
let kv: Vec<f32> = vec![1., 2., 3., 4., 5., 6., 7., 8.];
|
||
let out = expand_kv_heads(&kv, &cfg, 1, 1);
|
||
// Expect first 4 Q heads to carry KV[0] and next 4 to carry KV[1].
|
||
for q in 0..4 {
|
||
let slice = &out[q * 4..(q + 1) * 4];
|
||
assert_eq!(
|
||
slice,
|
||
&[1., 2., 3., 4.],
|
||
"Q head {q} should match KV head 0"
|
||
);
|
||
}
|
||
for q in 4..8 {
|
||
let slice = &out[q * 4..(q + 1) * 4];
|
||
assert_eq!(
|
||
slice,
|
||
&[5., 6., 7., 8.],
|
||
"Q head {q} should match KV head 1"
|
||
);
|
||
}
|
||
}
|
||
|
||
// ── gqa_attention_cpu output shape ────────────────────────────────────────
|
||
|
||
/// Output shape must be [num_q_heads, seq_len_q, head_dim].
|
||
#[test]
|
||
fn test_gqa_attention_output_shape() {
|
||
let cfg = GqaConfig::new(4, 2, 8).unwrap();
|
||
let seq_q = 3;
|
||
let seq_kv = 5;
|
||
let scale = (8f32).sqrt().recip();
|
||
let q: Vec<f32> = vec![0.1; 4 * seq_q * 8];
|
||
let k: Vec<f32> = vec![0.1; 2 * seq_kv * 8];
|
||
let v: Vec<f32> = vec![0.2; 2 * seq_kv * 8];
|
||
let out = gqa_attention_cpu(&q, &k, &v, &cfg, seq_q, seq_kv, scale, false);
|
||
assert_eq!(out.len(), 4 * seq_q * 8);
|
||
}
|
||
|
||
/// When num_kv == num_q, gqa_attention_cpu must match naive MHA within 1e-5.
|
||
#[test]
|
||
fn test_gqa_attention_mha_matches_standard() {
|
||
let n_heads = 2;
|
||
let seq = 3;
|
||
let head_dim = 4;
|
||
let cfg = GqaConfig::new(n_heads, n_heads, head_dim).unwrap();
|
||
let scale = (head_dim as f32).sqrt().recip();
|
||
|
||
// Simple ascending values for reproducibility.
|
||
let q: Vec<f32> = (0..(n_heads * seq * head_dim))
|
||
.map(|x| x as f32 * 0.05)
|
||
.collect();
|
||
let k = q.clone();
|
||
let v = q.clone();
|
||
|
||
let gqa_out = gqa_attention_cpu(&q, &k, &v, &cfg, seq, seq, scale, false);
|
||
|
||
// Compute naive MHA reference (same algorithm, but explicit expansion).
|
||
let ref_out = naive_mha_reference(&q, &k, &v, n_heads, seq, head_dim, scale, false);
|
||
|
||
assert_eq!(gqa_out.len(), ref_out.len());
|
||
for (a, b) in gqa_out.iter().zip(ref_out.iter()) {
|
||
assert!((a - b).abs() < 1e-5, "GQA/MHA mismatch: {a} vs {b}");
|
||
}
|
||
}
|
||
|
||
/// Causal mask: for seq=4, the output should reflect that future tokens are
|
||
/// invisible. Specifically, position 0 may only attend to key 0, so the
|
||
/// output for position 0 should equal `value[0]` exactly (after softmax of
|
||
/// a single-element distribution = 1.0).
|
||
#[test]
|
||
fn test_gqa_attention_causal_mask() {
|
||
let n_heads = 1;
|
||
let seq = 4;
|
||
let head_dim = 4;
|
||
let cfg = GqaConfig::new(n_heads, n_heads, head_dim).unwrap();
|
||
let scale = (head_dim as f32).sqrt().recip();
|
||
|
||
// Distinct value vectors so we can tell which positions contributed.
|
||
// V row j = all j+1 (1.0, 2.0, 3.0, 4.0 for j=0,1,2,3).
|
||
let q: Vec<f32> = vec![0.1; n_heads * seq * head_dim];
|
||
let k: Vec<f32> = vec![0.0; n_heads * seq * head_dim];
|
||
let mut v: Vec<f32> = vec![0.0; n_heads * seq * head_dim];
|
||
for j in 0..seq {
|
||
for d in 0..head_dim {
|
||
v[j * head_dim + d] = (j + 1) as f32;
|
||
}
|
||
}
|
||
|
||
let out = gqa_attention_cpu(&q, &k, &v, &cfg, seq, seq, scale, true);
|
||
|
||
// Position 0 (q_i=0): can only attend to k_j=0 → attn weight = 1.0 → out = V[0] = [1,1,1,1]
|
||
let pos0 = &out[0..head_dim];
|
||
for &val in pos0 {
|
||
assert!(
|
||
(val - 1.0).abs() < 1e-5,
|
||
"causal pos0: expected 1.0 got {val}"
|
||
);
|
||
}
|
||
|
||
// Position 3 (q_i=3): attends to k_j ∈ {0,1,2,3} uniformly (all keys identical).
|
||
// Expected output = mean of V rows = (1+2+3+4)/4 = 2.5.
|
||
let pos3 = &out[3 * head_dim..4 * head_dim];
|
||
for &val in pos3 {
|
||
assert!(
|
||
(val - 2.5).abs() < 1e-4,
|
||
"causal pos3: expected 2.5 got {val}"
|
||
);
|
||
}
|
||
}
|
||
|
||
// ── Helper ────────────────────────────────────────────────────────────────
|
||
|
||
/// Naive MHA reference (full expansion then standard attention).
|
||
fn naive_mha_reference(
|
||
q: &[f32],
|
||
k: &[f32],
|
||
v: &[f32],
|
||
n_heads: usize,
|
||
seq: usize,
|
||
head_dim: usize,
|
||
scale: f32,
|
||
causal: bool,
|
||
) -> Vec<f32> {
|
||
let mut out = vec![0.0f32; n_heads * seq * head_dim];
|
||
let head_stride = seq * head_dim;
|
||
let mut scores = vec![0.0f32; seq];
|
||
|
||
for h in 0..n_heads {
|
||
let q_base = h * head_stride;
|
||
let k_base = h * head_stride;
|
||
let v_base = h * head_stride;
|
||
let o_base = h * head_stride;
|
||
|
||
for i in 0..seq {
|
||
let q_row = &q[q_base + i * head_dim..q_base + i * head_dim + head_dim];
|
||
let max_j = if causal { i + 1 } else { seq };
|
||
|
||
for j in 0..seq {
|
||
if j < max_j {
|
||
let k_row = &k[k_base + j * head_dim..k_base + j * head_dim + head_dim];
|
||
scores[j] = dot(q_row, k_row) * scale;
|
||
} else {
|
||
scores[j] = f32::NEG_INFINITY;
|
||
}
|
||
}
|
||
|
||
let max_s = scores[..seq]
|
||
.iter()
|
||
.copied()
|
||
.fold(f32::NEG_INFINITY, f32::max);
|
||
let mut sum_e = 0.0f32;
|
||
for s in &mut scores[..seq] {
|
||
*s = (*s - max_s).exp();
|
||
sum_e += *s;
|
||
}
|
||
let inv = if sum_e > 0.0 { sum_e.recip() } else { 0.0 };
|
||
for s in &mut scores[..seq] {
|
||
*s *= inv;
|
||
}
|
||
|
||
let o_row = &mut out[o_base + i * head_dim..o_base + i * head_dim + head_dim];
|
||
for j in 0..seq {
|
||
let a = scores[j];
|
||
if a == 0.0 {
|
||
continue;
|
||
}
|
||
let v_row = &v[v_base + j * head_dim..v_base + j * head_dim + head_dim];
|
||
for d in 0..head_dim {
|
||
o_row[d] += a * v_row[d];
|
||
}
|
||
}
|
||
}
|
||
}
|
||
out
|
||
}
|
||
}
|