use std::path::Path;
use std::time::{Duration, Instant};
use rayon::prelude::*;
use anyhow::{Context as _, Result, ensure};
use crate::engine::gguf::{Gguf, TensorType};
use crate::engine::kernels::{
Q8Activation, dequantize_row, dim_to_f32, f32_to_fp16, fp16_to_f32, matrix_argmax,
matrix_matrix, matrix_matrix_pair, matrix_matrix_triple, matrix_vector, matrix_vector_triple,
matrix_vector_with_activation, rms_norm, softmax, vector_add,
};
use crate::engine::tokenizer::{Tokenizer, Utf8Decoder};
const LAYERS: usize = 24;
const HIDDEN: usize = 2_880;
const QUERY_HEADS: usize = 64;
const KEY_VALUE_HEADS: usize = 8;
const HEAD_DIMENSION: usize = 64;
const QUERY_SIZE: usize = QUERY_HEADS * HEAD_DIMENSION;
const KEY_VALUE_SIZE: usize = KEY_VALUE_HEADS * HEAD_DIMENSION;
const QUERY_GROUP: usize = QUERY_HEADS / KEY_VALUE_HEADS;
const EXPERTS: usize = 32;
const ACTIVE_EXPERTS: usize = 4;
const EXPERT_HIDDEN: usize = 2_880;
const VOCABULARY: usize = 201_088;
const SLIDING_WINDOW: usize = 128;
const PREFILL_BATCH_SIZE: usize = 128;
const EPSILON: f32 = 1.0e-5;
const ROPE_BASE: f32 = 150_000.0;
const ROPE_FACTOR: f32 = 32.0;
const ROPE_ORIGINAL_CONTEXT: f32 = 4_096.0;
const ROPE_BETA_FAST: f32 = 32.0;
const ROPE_BETA_SLOW: f32 = 1.0;
const SWIGLU_ALPHA: f32 = 1.702;
const SWIGLU_LIMIT: f32 = 7.0;
pub struct Generation {
pub text: String,
pub prompt_tokens: usize,
pub generated_tokens: usize,
pub prefill_duration: Duration,
pub decode_duration: Duration,
}
pub struct TextModel {
gguf: Gguf,
tokenizer: Tokenizer,
layers: Vec<LayerNames>,
}
struct LayerNames {
attention_norm: String,
query: String,
query_bias: String,
key: String,
key_bias: String,
value: String,
value_bias: String,
attention_output: String,
attention_output_bias: String,
attention_sinks: String,
feed_forward_norm: String,
router: String,
router_bias: String,
gate_experts: String,
gate_experts_bias: String,
up_experts: String,
up_experts_bias: String,
down_experts: String,
down_experts_bias: String,
}
#[derive(Clone, Copy)]
struct ExpertRoute {
token: usize,
slot: usize,
weight: f32,
}
impl LayerNames {
fn new(layer: usize) -> Self {
let prefix = format!("blk.{layer}");
Self {
attention_norm: format!("{prefix}.attn_norm.weight"),
query: format!("{prefix}.attn_q.weight"),
query_bias: format!("{prefix}.attn_q.bias"),
key: format!("{prefix}.attn_k.weight"),
key_bias: format!("{prefix}.attn_k.bias"),
value: format!("{prefix}.attn_v.weight"),
value_bias: format!("{prefix}.attn_v.bias"),
attention_output: format!("{prefix}.attn_output.weight"),
attention_output_bias: format!("{prefix}.attn_output.bias"),
attention_sinks: format!("{prefix}.attn_sinks.weight"),
feed_forward_norm: format!("{prefix}.post_attention_norm.weight"),
router: format!("{prefix}.ffn_gate_inp.weight"),
router_bias: format!("{prefix}.ffn_gate_inp.bias"),
gate_experts: format!("{prefix}.ffn_gate_exps.weight"),
gate_experts_bias: format!("{prefix}.ffn_gate_exps.bias"),
up_experts: format!("{prefix}.ffn_up_exps.weight"),
up_experts_bias: format!("{prefix}.ffn_up_exps.bias"),
down_experts: format!("{prefix}.ffn_down_exps.weight"),
down_experts_bias: format!("{prefix}.ffn_down_exps.bias"),
}
}
}
impl TextModel {
pub fn load(path: impl AsRef<Path>) -> Result<Self> {
let gguf = Gguf::load(path).context("load GPT-OSS-20B GGUF")?;
validate_metadata(&gguf).context("validate GPT-OSS-20B metadata")?;
validate_tensors(&gguf).context("validate GPT-OSS-20B tensors")?;
let tokenizer = Tokenizer::from_gpt_oss(&gguf).context("load GPT-OSS tokenizer")?;
Ok(Self {
gguf,
tokenizer,
layers: (0..LAYERS).map(LayerNames::new).collect(),
})
}
pub fn generate(
&self,
prompt: &str,
max_new_tokens: usize,
max_context_tokens: u16,
) -> Result<Generation> {
ensure!(max_context_tokens != 0, "context limit must not be zero");
let max_context_tokens = usize::from(max_context_tokens);
let prompt_ids = self
.tokenizer
.encode(prompt, true)
.context("tokenize prompt")?;
ensure!(
!prompt_ids.is_empty(),
"prompt must produce at least one token"
);
let requested = prompt_ids
.len()
.checked_add(max_new_tokens)
.context("requested token count overflow")?;
ensure!(
requested <= max_context_tokens,
"prompt ({}) plus generation ({max_new_tokens}) exceeds context limit {max_context_tokens}",
prompt_ids.len()
);
let prompt_tokens = prompt_ids.len();
let mut cache = KvCache::new(max_context_tokens)?;
let mut scratch = AttentionScratch::new(max_context_tokens)?;
let prefill_started = Instant::now();
let mut hidden = Vec::new();
for (batch, tokens) in prompt_ids.chunks(PREFILL_BATCH_SIZE).enumerate() {
let start = batch * PREFILL_BATCH_SIZE;
let positions = (start..start + tokens.len()).collect::<Vec<_>>();
let embeddings = self.embeddings(tokens)?;
hidden = self.forward_batch(&embeddings, &positions, &mut cache, &mut scratch)?;
}
let prefill_duration = prefill_started.elapsed();
if max_new_tokens == 0 {
return Ok(Generation {
text: String::new(),
prompt_tokens,
generated_tokens: 0,
prefill_duration,
decode_duration: Duration::ZERO,
});
}
let decode_started = Instant::now();
let last_hidden = &hidden[hidden.len() - HIDDEN..];
let mut token = self.greedy_token(last_hidden)?;
let mut decoder = Utf8Decoder::default();
let mut text = String::new();
let mut generated_tokens = 0;
while generated_tokens < max_new_tokens && !self.tokenizer.is_eos(token) {
generated_tokens += 1;
if let Some(piece) = decoder.push(self.tokenizer.piece(token)?)? {
text.push_str(&piece);
}
if generated_tokens == max_new_tokens {
break;
}
let position = prompt_tokens + generated_tokens - 1;
let embedding = self.embedding(token)?;
let hidden = self.forward(&embedding, position, &mut cache, &mut scratch)?;
token = self.greedy_token(&hidden)?;
}
text.push_str(&decoder.finish()?);
Ok(Generation {
text,
prompt_tokens,
generated_tokens,
prefill_duration,
decode_duration: decode_started.elapsed(),
})
}
fn embedding(&self, token: u32) -> Result<Vec<f32>> {
let tensor = self.gguf.tensor("token_embd.weight")?;
let mut embedding = vec![0.0; HIDDEN];
dequantize_row(tensor.q8_row(usize::try_from(token)?)?, &mut embedding)?;
Ok(embedding)
}
fn embeddings(&self, tokens: &[u32]) -> Result<Vec<f32>> {
let tensor = self.gguf.tensor("token_embd.weight")?;
let output_len = tokens
.len()
.checked_mul(HIDDEN)
.context("embedding batch size overflow")?;
let mut embeddings = vec![0.0; output_len];
for (token, embedding) in tokens.iter().zip(embeddings.chunks_exact_mut(HIDDEN)) {
dequantize_row(tensor.q8_row(usize::try_from(*token)?)?, embedding)?;
}
Ok(embeddings)
}
fn forward(
&self,
input: &[f32],
position: usize,
cache: &mut KvCache,
scratch: &mut AttentionScratch,
) -> Result<Vec<f32>> {
ensure!(input.len() == HIDDEN, "decoder input width differs");
let rope = Rope::new(position)?;
let mut hidden = input.to_vec();
for (layer, names) in self.layers.iter().enumerate() {
let normalized = rms_norm(
&hidden,
HIDDEN,
self.gguf.tensor(&names.attention_norm)?.f32_slice()?,
EPSILON,
)?;
let (mut query, mut key, mut value) = matrix_vector_triple(
&self.gguf.tensor(&names.query)?,
&self.gguf.tensor(&names.key)?,
&self.gguf.tensor(&names.value)?,
&normalized,
)?;
add_bias(
&mut query,
self.gguf.tensor(&names.query_bias)?.f32_slice()?,
)?;
add_bias(&mut key, self.gguf.tensor(&names.key_bias)?.f32_slice()?)?;
add_bias(
&mut value,
self.gguf.tensor(&names.value_bias)?.f32_slice()?,
)?;
apply_rope(&mut query, &rope)?;
apply_rope(&mut key, &rope)?;
cache.layers[layer].append(&key, &value)?;
let sliding = layer % 2 == 0;
causal_gqa(
&query,
&cache.layers[layer],
cache.layers[layer].token_count(),
self.gguf.tensor(&names.attention_sinks)?.f32_slice()?,
sliding,
scratch,
)?;
let mut projected =
matrix_vector(&self.gguf.tensor(&names.attention_output)?, &scratch.output)?;
add_bias(
&mut projected,
self.gguf
.tensor(&names.attention_output_bias)?
.f32_slice()?,
)?;
vector_add(&mut hidden, &projected)?;
let normalized = rms_norm(
&hidden,
HIDDEN,
self.gguf.tensor(&names.feed_forward_norm)?.f32_slice()?,
EPSILON,
)?;
let mixture = self.mixture_of_experts(&normalized, names)?;
vector_add(&mut hidden, &mixture)?;
}
rms_norm(
&hidden,
HIDDEN,
self.gguf.tensor("output_norm.weight")?.f32_slice()?,
EPSILON,
)
}
fn forward_batch(
&self,
input: &[f32],
positions: &[usize],
cache: &mut KvCache,
scratch: &mut AttentionScratch,
) -> Result<Vec<f32>> {
ensure!(!positions.is_empty(), "decoder batch is empty");
ensure!(
input.len() == positions.len() * HIDDEN,
"decoder batch input shape differs"
);
if positions.len() == 1 {
return self.forward(input, positions[0], cache, scratch);
}
let ropes = positions
.iter()
.copied()
.map(Rope::new)
.collect::<Result<Vec<_>>>()?;
let row_count = positions.len();
let mut hidden = input.to_vec();
for (layer, names) in self.layers.iter().enumerate() {
let normalized = rms_norm(
&hidden,
HIDDEN,
self.gguf.tensor(&names.attention_norm)?.f32_slice()?,
EPSILON,
)?;
let (mut query, mut key, mut value) = matrix_matrix_triple(
&self.gguf.tensor(&names.query)?,
&self.gguf.tensor(&names.key)?,
&self.gguf.tensor(&names.value)?,
&normalized,
row_count,
)?;
add_bias_batch(
&mut query,
self.gguf.tensor(&names.query_bias)?.f32_slice()?,
)?;
add_bias_batch(&mut key, self.gguf.tensor(&names.key_bias)?.f32_slice()?)?;
add_bias_batch(
&mut value,
self.gguf.tensor(&names.value_bias)?.f32_slice()?,
)?;
for ((query, key), rope) in query
.chunks_exact_mut(QUERY_SIZE)
.zip(key.chunks_exact_mut(KEY_VALUE_SIZE))
.zip(&ropes)
{
apply_rope(query, rope)?;
apply_rope(key, rope)?;
}
let cached_tokens = cache.layers[layer].token_count();
for (key, value) in key
.chunks_exact(KEY_VALUE_SIZE)
.zip(value.chunks_exact(KEY_VALUE_SIZE))
{
cache.layers[layer].append(key, value)?;
}
let sliding = layer % 2 == 0;
let sinks = self.gguf.tensor(&names.attention_sinks)?.f32_slice()?;
let mut attention = vec![0.0; row_count * QUERY_SIZE];
for (token, (query, output)) in query
.chunks_exact(QUERY_SIZE)
.zip(attention.chunks_exact_mut(QUERY_SIZE))
.enumerate()
{
causal_gqa(
query,
&cache.layers[layer],
cached_tokens + token + 1,
sinks,
sliding,
scratch,
)?;
output.copy_from_slice(&scratch.output);
}
let mut projected = matrix_matrix(
&self.gguf.tensor(&names.attention_output)?,
&attention,
row_count,
)?;
add_bias_batch(
&mut projected,
self.gguf
.tensor(&names.attention_output_bias)?
.f32_slice()?,
)?;
vector_add(&mut hidden, &projected)?;
let normalized = rms_norm(
&hidden,
HIDDEN,
self.gguf.tensor(&names.feed_forward_norm)?.f32_slice()?,
EPSILON,
)?;
let mixture = self.mixture_of_experts_batch(&normalized, names)?;
vector_add(&mut hidden, &mixture)?;
}
rms_norm(
&hidden,
HIDDEN,
self.gguf.tensor("output_norm.weight")?.f32_slice()?,
EPSILON,
)
}
fn mixture_of_experts(&self, input: &[f32], names: &LayerNames) -> Result<Vec<f32>> {
let input_activation = Q8Activation::new(input)?;
let mut router = matrix_vector_with_activation(
&self.gguf.tensor(&names.router)?,
input,
Some(&input_activation),
)?;
add_bias(
&mut router,
self.gguf.tensor(&names.router_bias)?.f32_slice()?,
)?;
let (indices, mut weights) = top_experts(&router);
softmax(&mut weights);
let gate_weights = self.gguf.tensor(&names.gate_experts)?;
let up_weights = self.gguf.tensor(&names.up_experts)?;
let down_weights = self.gguf.tensor(&names.down_experts)?;
let gate_biases = self.gguf.tensor(&names.gate_experts_bias)?;
let up_biases = self.gguf.tensor(&names.up_experts_bias)?;
let down_biases = self.gguf.tensor(&names.down_experts_bias)?;
let expert_outputs: Vec<Vec<f32>> = indices
.into_iter()
.zip(weights)
.collect::<Vec<_>>()
.into_par_iter()
.map(|(expert, routing_weight)| -> Result<Vec<f32>> {
let mut gate = matrix_vector_with_activation(
&gate_weights.matrix_slice(expert)?,
input,
Some(&input_activation),
)?;
let mut up = matrix_vector_with_activation(
&up_weights.matrix_slice(expert)?,
input,
Some(&input_activation),
)?;
add_bias(&mut gate, gate_biases.f32_row(expert)?)?;
add_bias(&mut up, up_biases.f32_row(expert)?)?;
for (gate, up) in gate.iter_mut().zip(&mut up) {
*gate = gate.min(SWIGLU_LIMIT);
*up = up.clamp(-SWIGLU_LIMIT, SWIGLU_LIMIT);
let glu = *gate / (1.0 + (-SWIGLU_ALPHA * *gate).exp());
*gate = (*up + 1.0) * glu;
}
let mut output = matrix_vector(&down_weights.matrix_slice(expert)?, &gate)?;
add_bias(&mut output, down_biases.f32_row(expert)?)?;
for value in &mut output {
*value *= routing_weight;
}
Ok(output)
})
.collect::<Result<_>>()?;
let mut mixture = vec![0.0; HIDDEN];
for output in expert_outputs {
for (mixed, value) in mixture.iter_mut().zip(output) {
*mixed += value;
}
}
Ok(mixture)
}
fn mixture_of_experts_batch(&self, input: &[f32], names: &LayerNames) -> Result<Vec<f32>> {
ensure!(
!input.is_empty() && input.len().is_multiple_of(HIDDEN),
"expert batch input shape differs"
);
let row_count = input.len() / HIDDEN;
let mut router = matrix_matrix(&self.gguf.tensor(&names.router)?, input, row_count)?;
add_bias_batch(
&mut router,
self.gguf.tensor(&names.router_bias)?.f32_slice()?,
)?;
let mut routes_by_expert = vec![Vec::<ExpertRoute>::new(); EXPERTS];
for (token, logits) in router.chunks_exact(EXPERTS).enumerate() {
let (indices, mut weights) = top_experts(logits);
softmax(&mut weights);
for (slot, (expert, weight)) in indices.into_iter().zip(weights).enumerate() {
routes_by_expert[expert].push(ExpertRoute {
token,
slot,
weight,
});
}
}
let gate_weights = self.gguf.tensor(&names.gate_experts)?;
let up_weights = self.gguf.tensor(&names.up_experts)?;
let down_weights = self.gguf.tensor(&names.down_experts)?;
let gate_biases = self.gguf.tensor(&names.gate_experts_bias)?;
let up_biases = self.gguf.tensor(&names.up_experts_bias)?;
let down_biases = self.gguf.tensor(&names.down_experts_bias)?;
let expert_outputs = routes_by_expert
.par_iter()
.enumerate()
.map(|(expert, expert_routes)| -> Result<Option<Vec<f32>>> {
if expert_routes.is_empty() {
return Ok(None);
}
let mut expert_input = Vec::with_capacity(expert_routes.len() * HIDDEN);
for route in expert_routes {
let start = route.token * HIDDEN;
expert_input.extend_from_slice(&input[start..start + HIDDEN]);
}
let (mut gate, mut up) = matrix_matrix_pair(
&gate_weights.matrix_slice(expert)?,
&up_weights.matrix_slice(expert)?,
&expert_input,
expert_routes.len(),
)?;
add_bias_batch(&mut gate, gate_biases.f32_row(expert)?)?;
add_bias_batch(&mut up, up_biases.f32_row(expert)?)?;
for (gate, up) in gate.iter_mut().zip(&mut up) {
*gate = gate.min(SWIGLU_LIMIT);
*up = up.clamp(-SWIGLU_LIMIT, SWIGLU_LIMIT);
let glu = *gate / (1.0 + (-SWIGLU_ALPHA * *gate).exp());
*gate = (*up + 1.0) * glu;
}
let mut output = matrix_matrix(
&down_weights.matrix_slice(expert)?,
&gate,
expert_routes.len(),
)?;
add_bias_batch(&mut output, down_biases.f32_row(expert)?)?;
for (output, route) in output.chunks_exact_mut(HIDDEN).zip(expert_routes) {
for value in output {
*value *= route.weight;
}
}
Ok(Some(output))
})
.collect::<Result<Vec<_>>>()?;
let route_values = row_count
.checked_mul(ACTIVE_EXPERTS)
.and_then(|routes| routes.checked_mul(HIDDEN))
.context("expert route workspace size overflow")?;
let mut route_outputs = vec![0.0; route_values];
for (expert_routes, output) in routes_by_expert.iter().zip(expert_outputs) {
let Some(output) = output else {
continue;
};
for (route, values) in expert_routes.iter().zip(output.chunks_exact(HIDDEN)) {
let start = (route.token * ACTIVE_EXPERTS + route.slot) * HIDDEN;
route_outputs[start..start + HIDDEN].copy_from_slice(values);
}
}
let mut mixture = vec![0.0; input.len()];
for (token, mixed) in mixture.chunks_exact_mut(HIDDEN).enumerate() {
for slot in 0..ACTIVE_EXPERTS {
let start = (token * ACTIVE_EXPERTS + slot) * HIDDEN;
for (value, routed) in mixed.iter_mut().zip(&route_outputs[start..start + HIDDEN]) {
*value += routed;
}
}
}
Ok(mixture)
}
fn greedy_token(&self, hidden: &[f32]) -> Result<u32> {
let output = self.gguf.tensor("output.weight")?;
let token = matrix_argmax(&output, hidden).context("compute greedy token")?;
u32::try_from(token).context("token ID exceeds u32")
}
}
fn add_bias(values: &mut [f32], bias: &[f32]) -> Result<()> {
ensure!(values.len() == bias.len(), "bias width differs");
values
.iter_mut()
.zip(bias)
.for_each(|(value, bias)| *value += bias);
Ok(())
}
fn add_bias_batch(values: &mut [f32], bias: &[f32]) -> Result<()> {
ensure!(
!bias.is_empty() && values.len().is_multiple_of(bias.len()),
"bias batch shape differs"
);
for row in values.chunks_exact_mut(bias.len()) {
for (value, bias) in row.iter_mut().zip(bias) {
*value += bias;
}
}
Ok(())
}
fn top_experts(logits: &[f32]) -> ([usize; ACTIVE_EXPERTS], [f32; ACTIVE_EXPERTS]) {
let mut indices = [0; ACTIVE_EXPERTS];
let mut values = [f32::NEG_INFINITY; ACTIVE_EXPERTS];
for (index, value) in logits.iter().copied().enumerate() {
let position = values
.iter()
.position(|selected| value > *selected)
.unwrap_or(ACTIVE_EXPERTS);
if position < ACTIVE_EXPERTS {
values[position..].rotate_right(1);
indices[position..].rotate_right(1);
values[position] = value;
indices[position] = index;
}
}
(indices, values)
}
struct Rope {
cosine: [f32; HEAD_DIMENSION / 2],
sine: [f32; HEAD_DIMENSION / 2],
}
impl Rope {
fn new(position: usize) -> Result<Self> {
let position = token_position_to_f32(position)?;
let attention_factor = 0.1 * ROPE_FACTOR.ln() + 1.0;
let low = correction_dimension(ROPE_BETA_FAST).max(0.0);
let high = correction_dimension(ROPE_BETA_SLOW).min(dim_to_f32(HEAD_DIMENSION - 1));
let mut rope = Self {
cosine: [0.0; HEAD_DIMENSION / 2],
sine: [0.0; HEAD_DIMENSION / 2],
};
for pair in 0..HEAD_DIMENSION / 2 {
let base_frequency =
ROPE_BASE.powf(-(dim_to_f32(2 * pair) / dim_to_f32(HEAD_DIMENSION)));
let ramp = ((dim_to_f32(pair) - low) / (high - low)).clamp(0.0, 1.0);
let frequency = base_frequency * (1.0 - ramp + ramp / ROPE_FACTOR);
let angle = position * frequency;
rope.cosine[pair] = angle.cos() * attention_factor;
rope.sine[pair] = angle.sin() * attention_factor;
}
Ok(rope)
}
}
fn token_position_to_f32(position: usize) -> Result<f32> {
Ok(f32::from(
u16::try_from(position).context("token position exceeds the supported context")?,
))
}
fn correction_dimension(rotations: f32) -> f32 {
dim_to_f32(HEAD_DIMENSION)
* (ROPE_ORIGINAL_CONTEXT / (rotations * 2.0 * std::f32::consts::PI)).ln()
/ (2.0 * ROPE_BASE.ln())
}
fn apply_rope(values: &mut [f32], rope: &Rope) -> Result<()> {
ensure!(
values.len().is_multiple_of(HEAD_DIMENSION),
"RoPE input width differs"
);
for head in values.chunks_exact_mut(HEAD_DIMENSION) {
let (first, second) = head.split_at_mut(HEAD_DIMENSION / 2);
for pair in 0..HEAD_DIMENSION / 2 {
let left = first[pair];
let right = second[pair];
first[pair] = left * rope.cosine[pair] - right * rope.sine[pair];
second[pair] = left * rope.sine[pair] + right * rope.cosine[pair];
}
}
Ok(())
}
#[derive(Default)]
struct LayerCache {
keys: Vec<u16>,
values: Vec<u16>,
}
impl LayerCache {
fn reserve(&mut self, tokens: usize) -> Result<()> {
let values = tokens
.checked_mul(KEY_VALUE_SIZE)
.context("key/value cache size overflow")?;
self.keys
.try_reserve_exact(values)
.context("reserve key cache")?;
self.values
.try_reserve_exact(values)
.context("reserve value cache")?;
Ok(())
}
fn append(&mut self, key: &[f32], value: &[f32]) -> Result<()> {
ensure!(
key.len() == KEY_VALUE_SIZE && value.len() == KEY_VALUE_SIZE,
"key/value width differs"
);
self.keys.extend(key.iter().copied().map(f32_to_fp16));
self.values.extend(value.iter().copied().map(f32_to_fp16));
Ok(())
}
fn token_count(&self) -> usize {
self.keys.len() / KEY_VALUE_SIZE
}
}
struct KvCache {
layers: Vec<LayerCache>,
}
impl KvCache {
fn new(tokens: usize) -> Result<Self> {
let mut layers = (0..LAYERS)
.map(|_| LayerCache::default())
.collect::<Vec<_>>();
for layer in &mut layers {
layer.reserve(tokens)?;
}
Ok(Self { layers })
}
}
#[derive(Default)]
struct AttentionScratch {
output: Vec<f32>,
scores: Vec<f32>,
}
impl AttentionScratch {
fn new(tokens: usize) -> Result<Self> {
let score_capacity = QUERY_HEADS
.checked_mul(tokens + 1)
.context("attention workspace overflow")?;
let mut scores = Vec::new();
scores
.try_reserve_exact(score_capacity)
.context("reserve attention workspace")?;
scores.resize(score_capacity, 0.0);
Ok(Self {
output: vec![0.0; QUERY_SIZE],
scores,
})
}
}
fn causal_gqa(
query: &[f32],
cache: &LayerCache,
token_count: usize,
sinks: &[f32],
sliding: bool,
scratch: &mut AttentionScratch,
) -> Result<()> {
ensure!(query.len() == QUERY_SIZE, "attention query width differs");
ensure!(sinks.len() == QUERY_HEADS, "attention sink count differs");
ensure!(
cache.keys.len() == cache.values.len(),
"key/value cache lengths differ"
);
let cached_tokens = cache.token_count();
ensure!(
token_count != 0 && token_count <= cached_tokens,
"invalid visible attention token count"
);
let first_token = if sliding {
token_count.saturating_sub(SLIDING_WINDOW)
} else {
0
};
let visible_tokens = token_count - first_token;
let score_width = visible_tokens + 1;
let scale = dim_to_f32(HEAD_DIMENSION).sqrt().recip();
scratch
.scores
.par_chunks_exact_mut(score_width)
.zip(scratch.output.par_chunks_exact_mut(HEAD_DIMENSION))
.zip(query.par_chunks_exact(HEAD_DIMENSION))
.enumerate()
.for_each(|(query_head, ((weights, output), query_values))| {
let key_value_head = query_head / QUERY_GROUP;
for (offset, weight) in weights[..visible_tokens].iter_mut().enumerate() {
let token = first_token + offset;
let start = token * KEY_VALUE_SIZE + key_value_head * HEAD_DIMENSION;
*weight = query_values
.iter()
.zip(&cache.keys[start..start + HEAD_DIMENSION])
.map(|(query, key)| query * fp16_to_f32(*key))
.sum::<f32>()
* scale;
}
weights[visible_tokens] = sinks[query_head];
softmax(weights);
output.fill(0.0);
for (offset, weight) in weights[..visible_tokens].iter().copied().enumerate() {
let token = first_token + offset;
let start = token * KEY_VALUE_SIZE + key_value_head * HEAD_DIMENSION;
for (channel, value) in output.iter_mut().enumerate() {
*value += weight * fp16_to_f32(cache.values[start + channel]);
}
}
});
ensure!(
scratch.output.iter().all(|value| value.is_finite()),
"attention output is not finite"
);
Ok(())
}
fn validate_metadata(gguf: &Gguf) -> Result<()> {
ensure!(
gguf.architecture() == "gpt-oss",
"expected gpt-oss architecture"
);
validate_u32(gguf, "gpt-oss.block_count", LAYERS)?;
validate_u32(gguf, "gpt-oss.embedding_length", HIDDEN)?;
validate_u32(gguf, "gpt-oss.feed_forward_length", EXPERT_HIDDEN)?;
validate_u32(gguf, "gpt-oss.attention.head_count", QUERY_HEADS)?;
validate_u32(gguf, "gpt-oss.attention.head_count_kv", KEY_VALUE_HEADS)?;
validate_u32(gguf, "gpt-oss.attention.key_length", HEAD_DIMENSION)?;
validate_u32(gguf, "gpt-oss.attention.value_length", HEAD_DIMENSION)?;
validate_u32(gguf, "gpt-oss.attention.sliding_window", SLIDING_WINDOW)?;
validate_u32(gguf, "gpt-oss.expert_count", EXPERTS)?;
validate_u32(gguf, "gpt-oss.expert_used_count", ACTIVE_EXPERTS)?;
validate_u32(gguf, "gpt-oss.expert_feed_forward_length", EXPERT_HIDDEN)?;
let epsilon = gguf.f32("gpt-oss.attention.layer_norm_rms_epsilon")?;
ensure!(
epsilon.to_bits() == EPSILON.to_bits(),
"RMS epsilon is {epsilon}, expected {EPSILON}"
);
let rope_base = gguf.f32("gpt-oss.rope.freq_base")?;
ensure!(
rope_base.to_bits() == ROPE_BASE.to_bits(),
"RoPE base is {rope_base}, expected {ROPE_BASE}"
);
ensure!(
gguf.strings("tokenizer.ggml.tokens")?.len() == VOCABULARY,
"vocabulary size differs"
);
ensure!(
gguf.string("tokenizer.ggml.model")? == "gpt2",
"expected GPT-2 tokenizer model"
);
ensure!(
gguf.string("tokenizer.ggml.pre")? == "gpt-4o",
"expected GPT-4o pre-tokenizer"
);
Ok(())
}
fn validate_u32(gguf: &Gguf, key: &str, expected: usize) -> Result<()> {
let value = gguf.u32(key)?;
ensure!(
value as usize == expected,
"{key} is {value}, expected {expected}"
);
Ok(())
}
fn validate_tensors(gguf: &Gguf) -> Result<()> {
validate_kind(
gguf,
"token_embd.weight",
&[HIDDEN, VOCABULARY],
TensorType::Q8_0,
)?;
validate_kind(gguf, "output_norm.weight", &[HIDDEN], TensorType::F32)?;
validate_kind(
gguf,
"output.weight",
&[HIDDEN, VOCABULARY],
TensorType::Q8_0,
)?;
for layer in 0..LAYERS {
let names = LayerNames::new(layer);
validate_kind(gguf, &names.attention_norm, &[HIDDEN], TensorType::F32)?;
validate_kind(gguf, &names.query, &[HIDDEN, QUERY_SIZE], TensorType::Q8_0)?;
validate_kind(gguf, &names.query_bias, &[QUERY_SIZE], TensorType::F32)?;
validate_kind(
gguf,
&names.key,
&[HIDDEN, KEY_VALUE_SIZE],
TensorType::Q8_0,
)?;
validate_kind(gguf, &names.key_bias, &[KEY_VALUE_SIZE], TensorType::F32)?;
validate_kind(
gguf,
&names.value,
&[HIDDEN, KEY_VALUE_SIZE],
TensorType::Q8_0,
)?;
validate_kind(gguf, &names.value_bias, &[KEY_VALUE_SIZE], TensorType::F32)?;
validate_kind(
gguf,
&names.attention_output,
&[QUERY_SIZE, HIDDEN],
TensorType::Q8_0,
)?;
validate_kind(
gguf,
&names.attention_output_bias,
&[HIDDEN],
TensorType::F32,
)?;
validate_kind(
gguf,
&names.attention_sinks,
&[QUERY_HEADS],
TensorType::F32,
)?;
validate_kind(gguf, &names.feed_forward_norm, &[HIDDEN], TensorType::F32)?;
validate_kind(gguf, &names.router, &[HIDDEN, EXPERTS], TensorType::F32)?;
validate_kind(gguf, &names.router_bias, &[EXPERTS], TensorType::F32)?;
for (weight, bias) in [
(&names.gate_experts, &names.gate_experts_bias),
(&names.up_experts, &names.up_experts_bias),
(&names.down_experts, &names.down_experts_bias),
] {
validate_kind(
gguf,
weight,
&[EXPERT_HIDDEN, HIDDEN, EXPERTS],
TensorType::Mxfp4,
)?;
validate_kind(gguf, bias, &[EXPERT_HIDDEN, EXPERTS], TensorType::F32)?;
}
}
Ok(())
}
fn validate_kind(gguf: &Gguf, name: &str, dimensions: &[usize], kind: TensorType) -> Result<()> {
let tensor = gguf.tensor(name)?;
ensure!(
tensor.dimensions() == dimensions,
"tensor `{name}` has dimensions {:?}, expected {dimensions:?}",
tensor.dimensions()
);
ensure!(
tensor.tensor_type() == kind,
"tensor `{name}` is {:?}, expected {kind:?}",
tensor.tensor_type()
);
Ok(())
}