use std::path::Path;
use std::time::{Duration, Instant};
use rayon::prelude::*;
use anyhow::{Context as _, Result, ensure};
use crate::engine::gguf::{Gguf, TensorHandle, TensorType};
use crate::engine::kernels::{
Q8Activation, accumulate_fp16, dequantize_row, dim_to_f32, dot_f32_fp16, matrix_argmax,
matrix_matrix, matrix_matrix_pair, matrix_matrix_triple, matrix_vector, matrix_vector_pair,
matrix_vector_triple, matrix_vector_with_activation, rms_norm, softmax, swiglu_inplace,
vector_add,
};
use crate::engine::tokenizer::{Tokenizer, Utf8Decoder};
use crate::engine::types::Fp16;
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<LayerWeights>,
token_embd: TensorHandle,
output_norm: TensorHandle,
output: TensorHandle,
}
struct LayerWeights {
attention_norm: TensorHandle,
query: TensorHandle,
query_bias: TensorHandle,
key: TensorHandle,
key_bias: TensorHandle,
value: TensorHandle,
value_bias: TensorHandle,
attention_output: TensorHandle,
attention_output_bias: TensorHandle,
attention_sinks: TensorHandle,
feed_forward_norm: TensorHandle,
router: TensorHandle,
router_bias: TensorHandle,
gate_experts: TensorHandle,
gate_experts_bias: TensorHandle,
up_experts: TensorHandle,
up_experts_bias: TensorHandle,
down_experts: TensorHandle,
down_experts_bias: TensorHandle,
}
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 LayerWeights {
fn resolve(gguf: &Gguf, names: &LayerNames) -> Result<Self> {
Ok(Self {
attention_norm: gguf.resolve(&names.attention_norm)?,
query: gguf.resolve(&names.query)?,
query_bias: gguf.resolve(&names.query_bias)?,
key: gguf.resolve(&names.key)?,
key_bias: gguf.resolve(&names.key_bias)?,
value: gguf.resolve(&names.value)?,
value_bias: gguf.resolve(&names.value_bias)?,
attention_output: gguf.resolve(&names.attention_output)?,
attention_output_bias: gguf.resolve(&names.attention_output_bias)?,
attention_sinks: gguf.resolve(&names.attention_sinks)?,
feed_forward_norm: gguf.resolve(&names.feed_forward_norm)?,
router: gguf.resolve(&names.router)?,
router_bias: gguf.resolve(&names.router_bias)?,
gate_experts: gguf.resolve(&names.gate_experts)?,
gate_experts_bias: gguf.resolve(&names.gate_experts_bias)?,
up_experts: gguf.resolve(&names.up_experts)?,
up_experts_bias: gguf.resolve(&names.up_experts_bias)?,
down_experts: gguf.resolve(&names.down_experts)?,
down_experts_bias: gguf.resolve(&names.down_experts_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")?;
let layers = (0..LAYERS)
.map(|layer| {
let names = LayerNames::new(layer);
LayerWeights::resolve(&gguf, &names)
})
.collect::<Result<Vec<_>>>()
.context("resolve layer weights")?;
let token_embd = gguf
.resolve("token_embd.weight")
.context("resolve token embeddings")?;
let output_norm = gguf
.resolve("output_norm.weight")
.context("resolve output norm")?;
let output = gguf
.resolve("output.weight")
.context("resolve output weights")?;
Ok(Self {
gguf,
tokenizer,
layers,
token_embd,
output_norm,
output,
})
}
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_from_handle(&self.token_embd);
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_from_handle(&self.token_embd);
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, weights) in self.layers.iter().enumerate() {
let attention_norm = self.gguf.tensor_from_handle(&weights.attention_norm);
let query_w = self.gguf.tensor_from_handle(&weights.query);
let key_w = self.gguf.tensor_from_handle(&weights.key);
let value_w = self.gguf.tensor_from_handle(&weights.value);
let query_bias = self.gguf.tensor_from_handle(&weights.query_bias);
let key_bias = self.gguf.tensor_from_handle(&weights.key_bias);
let value_bias = self.gguf.tensor_from_handle(&weights.value_bias);
let sinks = self.gguf.tensor_from_handle(&weights.attention_sinks);
let attention_output = self.gguf.tensor_from_handle(&weights.attention_output);
let attention_output_bias =
self.gguf.tensor_from_handle(&weights.attention_output_bias);
let feed_forward_norm = self.gguf.tensor_from_handle(&weights.feed_forward_norm);
let normalized = rms_norm(&hidden, HIDDEN, attention_norm.f32_slice()?, EPSILON)?;
let (mut query, mut key, mut value) =
matrix_vector_triple(&query_w, &key_w, &value_w, &normalized)?;
add_bias(&mut query, query_bias.f32_slice()?)?;
add_bias(&mut key, key_bias.f32_slice()?)?;
add_bias(&mut value, 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(),
sinks.f32_slice()?,
sliding,
scratch,
)?;
let mut projected = matrix_vector(&attention_output, scratch.output_rows(1))?;
add_bias(&mut projected, attention_output_bias.f32_slice()?)?;
vector_add(&mut hidden, &projected)?;
let normalized = rms_norm(&hidden, HIDDEN, feed_forward_norm.f32_slice()?, EPSILON)?;
let mixture = self.mixture_of_experts(&normalized, weights)?;
vector_add(&mut hidden, &mixture)?;
}
let output_norm = self.gguf.tensor_from_handle(&self.output_norm);
rms_norm(&hidden, HIDDEN, output_norm.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, weights) in self.layers.iter().enumerate() {
let attention_norm = self.gguf.tensor_from_handle(&weights.attention_norm);
let query_w = self.gguf.tensor_from_handle(&weights.query);
let key_w = self.gguf.tensor_from_handle(&weights.key);
let value_w = self.gguf.tensor_from_handle(&weights.value);
let query_bias = self.gguf.tensor_from_handle(&weights.query_bias);
let key_bias = self.gguf.tensor_from_handle(&weights.key_bias);
let value_bias = self.gguf.tensor_from_handle(&weights.value_bias);
let sinks = self.gguf.tensor_from_handle(&weights.attention_sinks);
let attention_output = self.gguf.tensor_from_handle(&weights.attention_output);
let attention_output_bias =
self.gguf.tensor_from_handle(&weights.attention_output_bias);
let feed_forward_norm = self.gguf.tensor_from_handle(&weights.feed_forward_norm);
let normalized = rms_norm(&hidden, HIDDEN, attention_norm.f32_slice()?, EPSILON)?;
let (mut query, mut key, mut value) =
matrix_matrix_triple(&query_w, &key_w, &value_w, &normalized, row_count)?;
add_bias_batch(&mut query, query_bias.f32_slice()?)?;
add_bias_batch(&mut key, key_bias.f32_slice()?)?;
add_bias_batch(&mut value, 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;
causal_gqa_batch(
&query,
&cache.layers[layer],
cached_tokens,
row_count,
sinks.f32_slice()?,
sliding,
scratch,
)?;
let mut projected =
matrix_matrix(&attention_output, scratch.output_rows(row_count), row_count)?;
add_bias_batch(&mut projected, attention_output_bias.f32_slice()?)?;
vector_add(&mut hidden, &projected)?;
let normalized = rms_norm(&hidden, HIDDEN, feed_forward_norm.f32_slice()?, EPSILON)?;
let mixture = self.mixture_of_experts_batch(&normalized, weights)?;
vector_add(&mut hidden, &mixture)?;
}
let output_norm = self.gguf.tensor_from_handle(&self.output_norm);
rms_norm(&hidden, HIDDEN, output_norm.f32_slice()?, EPSILON)
}
fn mixture_of_experts(&self, input: &[f32], layer: &LayerWeights) -> Result<Vec<f32>> {
let input_activation = Q8Activation::new(input)?;
let router_w = self.gguf.tensor_from_handle(&layer.router);
let router_bias = self.gguf.tensor_from_handle(&layer.router_bias);
let mut router = matrix_vector_with_activation(&router_w, input, Some(&input_activation))?;
add_bias(&mut router, router_bias.f32_slice()?)?;
let (indices, mut weights) = top_experts(&router);
softmax(&mut weights);
let gate_weights = self.gguf.tensor_from_handle(&layer.gate_experts);
let up_weights = self.gguf.tensor_from_handle(&layer.up_experts);
let down_weights = self.gguf.tensor_from_handle(&layer.down_experts);
let gate_biases = self.gguf.tensor_from_handle(&layer.gate_experts_bias);
let up_biases = self.gguf.tensor_from_handle(&layer.up_experts_bias);
let down_biases = self.gguf.tensor_from_handle(&layer.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, mut up) = matrix_vector_pair(
&gate_weights.matrix_slice(expert)?,
&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)?)?;
swiglu_inplace(&mut gate, &up, SWIGLU_ALPHA, SWIGLU_LIMIT)?;
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], layer: &LayerWeights) -> 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 router_w = self.gguf.tensor_from_handle(&layer.router);
let router_bias = self.gguf.tensor_from_handle(&layer.router_bias);
let mut router = matrix_matrix(&router_w, input, row_count)?;
add_bias_batch(&mut router, 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_from_handle(&layer.gate_experts);
let up_weights = self.gguf.tensor_from_handle(&layer.up_experts);
let down_weights = self.gguf.tensor_from_handle(&layer.down_experts);
let gate_biases = self.gguf.tensor_from_handle(&layer.gate_experts_bias);
let up_biases = self.gguf.tensor_from_handle(&layer.up_experts_bias);
let down_biases = self.gguf.tensor_from_handle(&layer.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)?)?;
swiglu_inplace(&mut gate, &up, SWIGLU_ALPHA, SWIGLU_LIMIT)?;
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_from_handle(&self.output);
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<Fp16>,
values: Vec<Fp16>,
}
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(Fp16::from));
self.values.extend(value.iter().copied().map(Fp16::from));
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>,
score_plane: usize,
}
impl AttentionScratch {
fn new(tokens: usize) -> Result<Self> {
let score_plane = QUERY_HEADS
.checked_mul(
tokens
.checked_add(1)
.context("attention score width overflow")?,
)
.context("attention score plane overflow")?;
let score_capacity = PREFILL_BATCH_SIZE
.checked_mul(score_plane)
.context("attention workspace overflow")?;
let output_capacity = PREFILL_BATCH_SIZE
.checked_mul(QUERY_SIZE)
.context("attention output workspace overflow")?;
let mut scores = Vec::new();
scores
.try_reserve_exact(score_capacity)
.context("reserve attention scores")?;
scores.resize(score_capacity, 0.0);
let mut output = Vec::new();
output
.try_reserve_exact(output_capacity)
.context("reserve attention output")?;
output.resize(output_capacity, 0.0);
Ok(Self {
output,
scores,
score_plane,
})
}
fn output_rows(&self, row_count: usize) -> &[f32] {
&self.output[..row_count * QUERY_SIZE]
}
fn token_workspace(&mut self) -> (&mut [f32], &mut [f32]) {
(
&mut self.output[..QUERY_SIZE],
&mut self.scores[..self.score_plane],
)
}
fn batch_workspace(&mut self, row_count: usize) -> Result<(&mut [f32], &mut [f32])> {
let output_values = row_count
.checked_mul(QUERY_SIZE)
.context("attention batch output overflow")?;
let score_values = row_count
.checked_mul(self.score_plane)
.context("attention batch score overflow")?;
ensure!(
output_values <= self.output.len(),
"attention batch exceeds output workspace"
);
ensure!(
score_values <= self.scores.len(),
"attention batch exceeds score workspace"
);
Ok((
&mut self.output[..output_values],
&mut self.scores[..score_values],
))
}
}
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, visible_tokens) = attention_window(token_count, sliding);
let score_width = visible_tokens + 1;
let score_values = QUERY_HEADS
.checked_mul(score_width)
.context("attention score width overflow")?;
let (output, scores) = scratch.token_workspace();
ensure!(
scores.len() >= score_values,
"attention score workspace too small"
);
let scale = dim_to_f32(HEAD_DIMENSION).sqrt().recip();
let scores = &mut scores[..score_values];
let ctx = AttentionHeadCtx {
cache,
sinks,
first_token,
visible_tokens,
scale,
};
scores
.par_chunks_exact_mut(score_width)
.zip(output.par_chunks_exact_mut(HEAD_DIMENSION))
.zip(query.par_chunks_exact(HEAD_DIMENSION))
.enumerate()
.for_each(|(query_head, ((weights, output), query_values))| {
attention_head(query_head, query_values, &ctx, weights, output);
});
ensure!(
output.iter().all(|value| value.is_finite()),
"attention output is not finite"
);
Ok(())
}
fn causal_gqa_batch(
queries: &[f32],
cache: &LayerCache,
cached_tokens: usize,
row_count: usize,
sinks: &[f32],
sliding: bool,
scratch: &mut AttentionScratch,
) -> Result<()> {
ensure!(row_count != 0, "attention batch is empty");
ensure!(
queries.len() == row_count * QUERY_SIZE,
"attention batch query shape differs"
);
ensure!(sinks.len() == QUERY_HEADS, "attention sink count differs");
ensure!(
cache.keys.len() == cache.values.len(),
"key/value cache lengths differ"
);
let final_tokens = cached_tokens
.checked_add(row_count)
.context("attention batch token count overflow")?;
ensure!(
final_tokens <= cache.token_count(),
"invalid visible attention token count"
);
let (outputs, scores) = scratch.batch_workspace(row_count)?;
let score_plane = scores.len() / row_count;
ensure!(
score_plane != 0 && scores.len() == row_count * score_plane,
"attention batch score shape differs"
);
let scale = dim_to_f32(HEAD_DIMENSION).sqrt().recip();
outputs
.par_chunks_exact_mut(QUERY_SIZE)
.zip(queries.par_chunks_exact(QUERY_SIZE))
.zip(scores.par_chunks_exact_mut(score_plane))
.enumerate()
.for_each(|(token, ((output, query), score_plane))| {
let token_count = cached_tokens + token + 1;
let (first_token, visible_tokens) = attention_window(token_count, sliding);
let score_width = visible_tokens + 1;
let ctx = AttentionHeadCtx {
cache,
sinks,
first_token,
visible_tokens,
scale,
};
for query_head in 0..QUERY_HEADS {
let weights =
&mut score_plane[query_head * score_width..(query_head + 1) * score_width];
let head_output =
&mut output[query_head * HEAD_DIMENSION..(query_head + 1) * HEAD_DIMENSION];
let query_values =
&query[query_head * HEAD_DIMENSION..(query_head + 1) * HEAD_DIMENSION];
attention_head(query_head, query_values, &ctx, weights, head_output);
}
});
ensure!(
outputs.iter().all(|value| value.is_finite()),
"attention output is not finite"
);
Ok(())
}
fn attention_window(token_count: usize, sliding: bool) -> (usize, usize) {
let first_token = if sliding {
token_count.saturating_sub(SLIDING_WINDOW)
} else {
0
};
(first_token, token_count - first_token)
}
struct AttentionHeadCtx<'a> {
cache: &'a LayerCache,
sinks: &'a [f32],
first_token: usize,
visible_tokens: usize,
scale: f32,
}
fn attention_head(
query_head: usize,
query_values: &[f32],
ctx: &AttentionHeadCtx<'_>,
weights: &mut [f32],
output: &mut [f32],
) {
let key_value_head = query_head / QUERY_GROUP;
for (offset, weight) in weights[..ctx.visible_tokens].iter_mut().enumerate() {
let token = ctx.first_token + offset;
let start = token * KEY_VALUE_SIZE + key_value_head * HEAD_DIMENSION;
*weight =
dot_f32_fp16(query_values, &ctx.cache.keys[start..start + HEAD_DIMENSION]) * ctx.scale;
}
weights[ctx.visible_tokens] = ctx.sinks[query_head];
softmax(weights);
output.fill(0.0);
for (offset, weight) in weights[..ctx.visible_tokens].iter().copied().enumerate() {
let token = ctx.first_token + offset;
let start = token * KEY_VALUE_SIZE + key_value_head * HEAD_DIMENSION;
accumulate_fp16(
output,
&ctx.cache.values[start..start + HEAD_DIMENSION],
weight,
);
}
}
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(())
}