use models::layout::{AttentionOutput, DecoderConfig};
use super::{
Array, Error, KvCache, ModelTensors, NormWeight, PagedContextMode, QuantizedLinear, Result,
RopeOptions, Stream, fused_gate_up::split_last, native_paged_attention_mode,
};
#[derive(Debug, Clone, Copy)]
pub struct GatedFullAttentionConfig {
attention_heads: i32,
key_value_heads: i32,
head_dim: i32,
rope_dimensions: i32,
rope_base: f32,
attention_scale: f32,
rms_norm_eps: f32,
}
#[derive(Debug)]
pub struct GatedFullAttention {
config: GatedFullAttentionConfig,
query: QuantizedLinear,
key: QuantizedLinear,
value: QuantizedLinear,
output: QuantizedLinear,
query_norm: NormWeight,
key_norm: NormWeight,
}
impl GatedFullAttentionConfig {
pub fn from_decoder(decoder: &DecoderConfig) -> Result<Self> {
if decoder.attention_output != AttentionOutput::Gated {
return Err(Error::InvalidModel("gated full attention requires an output gate".into()));
}
let head_dim = i32::try_from(decoder.head_dim)?;
let partial = decoder.partial_rotary_factor.unwrap_or(1.0);
let rope_dimensions: i32 = (f64::from(head_dim) * partial).round().to_string().parse()?;
let config = Self {
attention_heads: i32::try_from(decoder.num_attention_heads)?,
key_value_heads: i32::try_from(decoder.num_key_value_heads)?,
head_dim,
rope_dimensions,
rope_base: decoder.rope_theta.unwrap_or(10_000.0).to_string().parse()?,
attention_scale: head_dim.to_string().parse::<f32>()?.sqrt().recip(),
rms_norm_eps: decoder.rms_norm_eps.to_string().parse()?,
};
config.validate()?;
Ok(config)
}
fn validate(self) -> Result<()> {
if self.attention_heads <= 0
|| self.key_value_heads <= 0
|| self.head_dim <= 0
|| self.attention_heads % self.key_value_heads != 0
|| self.rope_dimensions <= 0
|| self.rope_dimensions > self.head_dim
|| self.rope_dimensions % 2 != 0
|| !self.rope_base.is_finite()
|| !self.attention_scale.is_finite()
|| !self.rms_norm_eps.is_finite()
{
return Err(Error::InvalidModel(format!(
"invalid gated full attention configuration: {self:?}"
)));
}
Ok(())
}
}
impl GatedFullAttention {
pub fn load_with_norm_shift(
tensors: &ModelTensors,
prefix: &str,
config: GatedFullAttentionConfig,
group_size: i32,
norm_shift: f32,
stream: &Stream,
) -> Result<Self> {
Ok(Self {
config,
query: linear(tensors, prefix, "q_proj", group_size)?,
key: linear(tensors, prefix, "k_proj", group_size)?,
value: linear(tensors, prefix, "v_proj", group_size)?,
output: linear(tensors, prefix, "o_proj", group_size)?,
query_norm: NormWeight::load_adjusted(
tensors,
&format!("{prefix}.q_norm"),
norm_shift,
stream,
)?,
key_norm: NormWeight::load_adjusted(
tensors,
&format!("{prefix}.k_norm"),
norm_shift,
stream,
)?,
})
}
pub fn forward(
&self,
input: &Array,
cache: &mut KvCache,
page_min_context: usize,
position: i32,
causal: bool,
stream: &Stream,
) -> Result<Array> {
let shape = input.shape()?;
let batch = dimension(&shape, 0, "batch")?;
let sequence = dimension(&shape, 1, "sequence")?;
let heads = self.config.attention_heads;
let head_dim = self.config.head_dim;
let query_width = heads.checked_mul(head_dim).ok_or(Error::ShapeOverflow)?;
let projected = self.query.forward(input, stream)?.reshape(
&[batch, sequence, heads, head_dim.checked_mul(2).ok_or(Error::ShapeOverflow)?],
stream,
)?;
let (queries, gate) = split_last(&projected, usize::try_from(head_dim)?, stream)?;
let gate = gate.reshape(&[batch, sequence, query_width], stream)?;
let queries = self.query_norm.apply(&queries, self.config.rms_norm_eps, stream)?;
let queries = rope(&queries, self.config, position, stream)?;
let keys = self
.key
.forward(input, stream)?
.reshape(&[batch, sequence, self.config.key_value_heads, head_dim], stream)?;
let keys = self.key_norm.apply(&keys, self.config.rms_norm_eps, stream)?;
let keys = rope(&keys, self.config, position, stream)?;
let values = self
.value
.forward(input, stream)?
.reshape(&[batch, sequence, self.config.key_value_heads, head_dim], stream)?
.transpose(&[0, 2, 1, 3], stream)?;
let mode = if sequence == 1 {
native_paged_attention_mode(
head_dim,
heads,
self.config.key_value_heads,
usize::try_from(position)? + 1,
stream.config().cache.force_native_paged_attention,
)
} else {
PagedContextMode::View
};
let context =
cache.update_for_attention_mode(&keys, &values, stream, page_min_context, mode)?;
let attended = match context.paged {
Some(paged) => queries.paged_scaled_dot_product_attention_with_scratch(
paged.attention(),
paged.scratch(),
self.config.attention_scale,
stream,
)?,
None => queries.scaled_dot_product_attention(
&context.keys,
&context.values,
self.config.attention_scale,
causal,
stream,
)?,
};
let attended = attended
.transpose(&[0, 2, 1, 3], stream)?
.reshape(&[batch, sequence, query_width], stream)?;
self.output.forward(&gate.sigmoid_mul(&attended, stream)?, stream)
}
}
fn rope(
input: &Array,
config: GatedFullAttentionConfig,
position: i32,
stream: &Stream,
) -> Result<Array> {
input.transpose(&[0, 2, 1, 3], stream)?.rope(
RopeOptions {
dimensions: config.rope_dimensions,
traditional: false,
base: Some(config.rope_base),
scale: 1.0,
offset: position,
},
stream,
)
}
fn dimension(shape: &[i32], axis: usize, name: &str) -> Result<i32> {
shape
.get(axis)
.copied()
.ok_or_else(|| Error::InvalidModel(format!("attention input has no {name} axis")))
}
fn linear(
tensors: &ModelTensors,
prefix: &str,
name: &str,
group_size: i32,
) -> Result<QuantizedLinear> {
QuantizedLinear::load(tensors, &format!("{prefix}.{name}"), group_size)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::engine::QuantizedArrays;
#[test]
fn executes_gated_grouped_query_attention_on_the_gpu_stream() -> Result<()> {
let stream = Stream::new_gpu()?;
let config = GatedFullAttentionConfig {
attention_heads: 1,
key_value_heads: 1,
head_dim: 64,
rope_dimensions: 64,
rope_base: 10_000.0,
attention_scale: 0.125,
rms_norm_eps: 1.0e-6,
};
let attention = GatedFullAttention {
config,
query: linear(128, &stream)?,
key: linear(64, &stream)?,
value: linear(64, &stream)?,
output: linear(64, &stream)?,
query_norm: NormWeight::from_weight(Array::from_f32(&vec![1.0; 64], &[64])?),
key_norm: NormWeight::from_weight(Array::from_f32(&vec![1.0; 64], &[64])?),
};
let input = Array::from_f32(&vec![0.0; 128], &[1, 2, 64])?;
let mut cache = KvCache::new(16)?;
let output = attention.forward(&input, &mut cache, 0, 0, true, &stream)?;
output.async_eval()?;
stream.synchronize()?;
assert_eq!(output.shape()?, vec![1, 2, 64]);
assert!(output.to_vec_f32()?.iter().all(|value| *value == 0.0));
assert_eq!(cache.offset()?, 2);
Ok(())
}
fn linear(output_width: i32, stream: &Stream) -> Result<QuantizedLinear> {
let dense =
Array::from_f32(&vec![0.0; usize::try_from(output_width * 64)?], &[output_width, 64])?;
let arrays: QuantizedArrays = dense.quantize(64, 4, stream)?;
Ok(QuantizedLinear::from_quantized(arrays, 64, 4))
}
}