use std::env;
use std::path::Path;
use std::time::{Duration, Instant};
use super::expert_store::{ExpertKey, ExpertLease, ExpertRequest, ExpertStore};
use super::profile::{ProcessDelta, ProcessSnapshot};
use rayon::prelude::*;
use anyhow::{Context as _, Result, ensure};
use crate::engine::model::gguf::{ByteRange, Gguf, Tensor, TensorHandle, TensorType};
use crate::engine::model::kernels::{
Q8Activation, accumulate_fp16, dequantize_row, 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::model::tokenizer::{Tokenizer, Utf8Decoder};
use crate::engine::model::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 PREFILL_EXPERT_PARALLEL_ROWS: usize = 4 * PREFILL_BATCH_SIZE;
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;
const LOOP_HISTORY_TOKENS: usize = 128;
const LOOP_MAX_PERIOD: usize = 32;
const LOOP_MIN_REPETITIONS: usize = 4;
const LOOP_MIN_TOKENS: usize = 16;
const MEBIBYTE: u64 = 1024 * 1024;
const GIBIBYTE: u64 = 1024 * MEBIBYTE;
const DEFAULT_EXPERT_CACHE_BYTES: u64 = 2 * GIBIBYTE;
const MAX_EXPERT_CACHE_BYTES: u64 = 16 * GIBIBYTE;
const EXPERT_CACHE_SIZE_ENV: &str = "GOOSEDUMP_EXPERT_CACHE_SIZE";
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum GenerationStopReason {
EndOfSequence,
TokenLimit,
TokenCycle,
TimeBudget,
}
pub struct Generation {
pub text: String,
pub prompt_tokens: usize,
pub generated_tokens: usize,
pub decode_tokens: usize,
#[cfg(feature = "bench")]
pub generated_token_ids: Vec<u32>,
pub prefill_duration: Duration,
pub decode_duration: Duration,
pub stop_reason: GenerationStopReason,
pub(super) profile: Option<GenerationProfile>,
}
pub struct TextModel {
gguf: Gguf,
tokenizer: Tokenizer,
layers: Vec<LayerWeights>,
expert_store: ExpertStore,
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,
expert_bundles: Vec<ExpertBundle>,
}
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,
}
#[derive(Clone, Copy)]
struct ExpertBundle {
ranges: [ByteRange; 6],
}
struct ExpertWeights<'a> {
gate: Tensor<'a>,
gate_bias: Tensor<'a>,
up: Tensor<'a>,
up_bias: Tensor<'a>,
down: Tensor<'a>,
down_bias: Tensor<'a>,
}
impl ExpertBundle {
fn new(parent_ranges: [ByteRange; 6], expert: usize) -> Result<Self> {
let [gate, gate_bias, up, up_bias, down, down_bias] = parent_ranges;
Ok(Self {
ranges: [
gate.split(EXPERTS, expert)?,
gate_bias.split(EXPERTS, expert)?,
up.split(EXPERTS, expert)?,
up_bias.split(EXPERTS, expert)?,
down.split(EXPERTS, expert)?,
down_bias.split(EXPERTS, expert)?,
],
})
}
fn byte_len(self) -> Result<usize> {
self.ranges.iter().try_fold(0usize, |total, range| {
total
.checked_add(range.len())
.context("expert bundle size overflow")
})
}
}
#[derive(Default)]
pub(super) struct ExpertRouteProfile {
counts: [[u64; EXPERTS]; LAYERS],
}
impl ExpertRouteProfile {
fn record(&mut self, layer: usize, experts: &[usize; ACTIVE_EXPERTS]) {
let Some(counts) = self.counts.get_mut(layer) else {
return;
};
for &expert in experts {
if let Some(count) = counts.get_mut(expert) {
*count = count.saturating_add(1);
}
}
}
pub(super) fn selections(&self) -> u64 {
self.counts
.iter()
.flatten()
.copied()
.fold(0, u64::saturating_add)
}
pub(super) fn unique_layer_experts(&self) -> u64 {
self.counts
.iter()
.flatten()
.map(|count| u64::from(*count != 0))
.sum()
}
pub(super) fn layer_unique_counts(&self) -> Vec<usize> {
self.counts
.iter()
.map(|counts| counts.iter().filter(|count| **count != 0).count())
.collect()
}
pub(super) fn route_counts(&self) -> impl Iterator<Item = u64> + '_ {
self.counts.iter().flatten().copied()
}
}
pub(super) struct PhaseProfile {
pub(super) process: Option<ProcessDelta>,
pub(super) expert_routes: ExpertRouteProfile,
}
pub(super) struct GenerationProfile {
pub(super) prefill: PhaseProfile,
pub(super) decode: PhaseProfile,
}
struct DecodeOutput {
text: String,
generated_tokens: usize,
decode_tokens: usize,
#[cfg(feature = "bench")]
generated_token_ids: Vec<u32>,
duration: Duration,
stop_reason: GenerationStopReason,
profile: PhaseProfile,
}
fn process_snapshot(enabled: bool) -> Option<ProcessSnapshot> {
enabled.then(ProcessSnapshot::capture).flatten()
}
fn process_delta(enabled: bool, start: Option<ProcessSnapshot>) -> Option<ProcessDelta> {
ProcessSnapshot::delta(start, process_snapshot(enabled))
}
#[derive(Clone, Copy)]
struct DecodeConfig {
profile_enabled: bool,
max_duration: Option<Duration>,
}
struct DecodeGuard {
tokens: Vec<u32>,
max_duration: Option<Duration>,
}
impl DecodeGuard {
fn new(max_duration: Option<Duration>) -> Self {
Self {
tokens: Vec::with_capacity(LOOP_HISTORY_TOKENS),
max_duration,
}
}
fn time_budget_exhausted(&self, elapsed: Duration) -> bool {
self.max_duration
.is_some_and(|max_duration| elapsed >= max_duration)
}
fn observe_token(&mut self, token: u32) -> bool {
self.tokens.push(token);
if self.tokens.len() > LOOP_HISTORY_TOKENS {
self.tokens.remove(0);
}
repeated_token_cycle(&self.tokens)
}
}
fn repeated_token_cycle(tokens: &[u32]) -> bool {
let max_period = LOOP_MAX_PERIOD.min(tokens.len() / LOOP_MIN_REPETITIONS);
(1..=max_period).any(|period| {
let repetitions = LOOP_MIN_REPETITIONS.max(LOOP_MIN_TOKENS.div_ceil(period));
let span = period * repetitions;
if span > tokens.len() {
return false;
}
let suffix = &tokens[tokens.len() - span..];
let pattern = &suffix[..period];
suffix.chunks_exact(period).all(|chunk| chunk == pattern)
})
}
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> {
let gate_experts = gguf.resolve(&names.gate_experts)?;
let gate_experts_bias = gguf.resolve(&names.gate_experts_bias)?;
let up_experts = gguf.resolve(&names.up_experts)?;
let up_experts_bias = gguf.resolve(&names.up_experts_bias)?;
let down_experts = gguf.resolve(&names.down_experts)?;
let down_experts_bias = gguf.resolve(&names.down_experts_bias)?;
let expert_parent_ranges = [
gate_experts.byte_range(),
gate_experts_bias.byte_range(),
up_experts.byte_range(),
up_experts_bias.byte_range(),
down_experts.byte_range(),
down_experts_bias.byte_range(),
];
let expert_bundles = (0..EXPERTS)
.map(|expert| ExpertBundle::new(expert_parent_ranges, expert))
.collect::<Result<Vec<_>>>()?;
let weights = 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,
gate_experts_bias,
up_experts,
up_experts_bias,
down_experts,
down_experts_bias,
expert_bundles,
};
weights.validate_expert_bundles(expert_parent_ranges)?;
Ok(weights)
}
fn validate_expert_bundles(&self, parent_ranges: [ByteRange; 6]) -> Result<()> {
ensure!(
self.expert_bundles.len() == EXPERTS,
"expected {EXPERTS} expert bundles, found {}",
self.expert_bundles.len()
);
for (range_index, parent) in parent_ranges.into_iter().enumerate() {
let mut expected_start = parent.start();
for bundle in &self.expert_bundles {
let range = bundle.ranges[range_index];
ensure!(
range.start() == expected_start,
"expert bundle ranges are not contiguous"
);
expected_start = range.end();
}
ensure!(
expected_start == parent.end(),
"expert bundle ranges do not cover their tensor"
);
}
Ok(())
}
fn expert_weights<'a>(&'a self, lease: &'a ExpertLease) -> Result<ExpertWeights<'a>> {
Ok(ExpertWeights {
gate: self.gate_experts.partition_from_bytes(lease.part(0)?)?,
gate_bias: self
.gate_experts_bias
.partition_from_bytes(lease.part(1)?)?,
up: self.up_experts.partition_from_bytes(lease.part(2)?)?,
up_bias: self.up_experts_bias.partition_from_bytes(lease.part(3)?)?,
down: self.down_experts.partition_from_bytes(lease.part(4)?)?,
down_bias: self
.down_experts_bias
.partition_from_bytes(lease.part(5)?)?,
})
}
}
fn expert_slot_bytes(layers: &[LayerWeights]) -> Result<usize> {
let first_bundle = layers
.first()
.and_then(|layer| layer.expert_bundles.first())
.context("model has no expert bundles")?;
let slot_bytes = first_bundle.byte_len()?;
for bundle in layers.iter().flat_map(|layer| &layer.expert_bundles) {
ensure!(
bundle.byte_len()? == slot_bytes,
"expert bundles do not share one cache slot size"
);
}
Ok(slot_bytes)
}
fn parse_expert_cache_size(value: &str) -> Result<usize> {
let Some((suffix_index, suffix)) = value.char_indices().next_back() else {
anyhow::bail!("{EXPERT_CACHE_SIZE_ENV} must be an integer followed by M or G");
};
let amount = &value[..suffix_index];
ensure!(
!amount.is_empty() && amount.bytes().all(|byte| byte.is_ascii_digit()),
"{EXPERT_CACHE_SIZE_ENV} must be an integer followed by M or G"
);
let multiplier = match suffix {
'M' => MEBIBYTE,
'G' => GIBIBYTE,
_ => anyhow::bail!("{EXPERT_CACHE_SIZE_ENV} must be an integer followed by M or G"),
};
let bytes = amount
.parse::<u64>()
.with_context(|| format!("parse {EXPERT_CACHE_SIZE_ENV}"))?
.checked_mul(multiplier)
.context("expert cache byte budget overflow")?;
ensure!(
(MEBIBYTE..=MAX_EXPERT_CACHE_BYTES).contains(&bytes),
"{EXPERT_CACHE_SIZE_ENV} must be between 1M and 16G"
);
usize::try_from(bytes).context("expert cache size exceeds platform address space")
}
fn expert_cache_bytes() -> Result<usize> {
let Some(value) = env::var_os(EXPERT_CACHE_SIZE_ENV) else {
return usize::try_from(DEFAULT_EXPERT_CACHE_BYTES)
.context("default expert cache size exceeds platform address space");
};
let value = value
.into_string()
.map_err(|_| anyhow::anyhow!("{EXPERT_CACHE_SIZE_ENV} is not valid UTF-8"))?;
parse_expert_cache_size(&value)
}
impl TextModel {
pub fn load(path: impl AsRef<Path>) -> Result<Self> {
let path = path.as_ref();
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 expert_store = ExpertStore::new(
gguf.source_file()?,
expert_slot_bytes(&layers)?,
expert_cache_bytes()?,
EXPERTS,
LAYERS * EXPERTS,
)?;
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,
expert_store,
token_embd,
output_norm,
output,
})
}
pub fn prompt_tokens(&self, prompt: &str) -> Result<usize> {
self.tokenizer
.encode(prompt, true)
.map(|tokens| tokens.len())
.context("tokenize prompt")
}
pub fn generate(
&self,
prompt: &str,
max_new_tokens: usize,
max_context_tokens: u16,
profile_enabled: bool,
max_decode_duration: Option<Duration>,
) -> 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 mut prefill_routes = ExpertRouteProfile::default();
let prefill_process_start = process_snapshot(profile_enabled);
let prefill_started = Instant::now();
let hidden = self.prefill(
&prompt_ids,
&mut cache,
&mut scratch,
profile_enabled.then_some(&mut prefill_routes),
)?;
if max_new_tokens == 0 {
let prefill_duration = prefill_started.elapsed();
let prefill_process = process_delta(profile_enabled, prefill_process_start);
return Ok(Self::empty_generation(
prompt_tokens,
prefill_duration,
prefill_process,
prefill_routes,
profile_enabled,
));
}
let last_hidden_start = hidden
.len()
.checked_sub(HIDDEN)
.context("prefill output is shorter than the hidden width")?;
let last_hidden = hidden
.get(last_hidden_start..)
.context("prefill output range is invalid")?;
let token = self.greedy_token(last_hidden)?;
let prefill_duration = prefill_started.elapsed();
let prefill_process = process_delta(profile_enabled, prefill_process_start);
let decoded = self.decode(
token,
prompt_tokens,
max_new_tokens,
&mut cache,
&mut scratch,
DecodeConfig {
profile_enabled,
max_duration: max_decode_duration,
},
)?;
Ok(Generation {
text: decoded.text,
prompt_tokens,
generated_tokens: decoded.generated_tokens,
decode_tokens: decoded.decode_tokens,
#[cfg(feature = "bench")]
generated_token_ids: decoded.generated_token_ids,
prefill_duration,
decode_duration: decoded.duration,
stop_reason: decoded.stop_reason,
profile: profile_enabled.then_some(GenerationProfile {
prefill: PhaseProfile {
process: prefill_process,
expert_routes: prefill_routes,
},
decode: decoded.profile,
}),
})
}
fn empty_generation(
prompt_tokens: usize,
prefill_duration: Duration,
prefill_process: Option<ProcessDelta>,
prefill_routes: ExpertRouteProfile,
profile_enabled: bool,
) -> Generation {
Generation {
text: String::new(),
prompt_tokens,
generated_tokens: 0,
decode_tokens: 0,
#[cfg(feature = "bench")]
generated_token_ids: Vec::new(),
prefill_duration,
decode_duration: Duration::ZERO,
stop_reason: GenerationStopReason::TokenLimit,
profile: profile_enabled.then_some(GenerationProfile {
prefill: PhaseProfile {
process: prefill_process,
expert_routes: prefill_routes,
},
decode: PhaseProfile {
process: None,
expert_routes: ExpertRouteProfile::default(),
},
}),
}
}
fn decode(
&self,
mut token: u32,
prompt_tokens: usize,
max_new_tokens: usize,
cache: &mut KvCache,
scratch: &mut AttentionScratch,
config: DecodeConfig,
) -> Result<DecodeOutput> {
let DecodeConfig {
profile_enabled,
max_duration,
} = config;
let mut decoder = Utf8Decoder::default();
let mut text = String::new();
#[cfg(feature = "bench")]
let mut generated_token_ids = Vec::with_capacity(max_new_tokens);
let mut routes = ExpertRouteProfile::default();
let mut generated_tokens = 0;
let mut decode_tokens = 0;
let mut guard = DecodeGuard::new(max_duration);
let mut stop_reason = GenerationStopReason::EndOfSequence;
let process_start = process_snapshot(profile_enabled);
let started = Instant::now();
while generated_tokens < max_new_tokens && !self.tokenizer.is_eos(token) {
if guard.time_budget_exhausted(started.elapsed()) {
stop_reason = GenerationStopReason::TimeBudget;
break;
}
generated_tokens += 1;
#[cfg(feature = "bench")]
generated_token_ids.push(token);
if let Some(piece) = decoder.push(self.tokenizer.piece(token)?)? {
text.push_str(&piece);
}
if guard.observe_token(token) {
stop_reason = GenerationStopReason::TokenCycle;
break;
}
if generated_tokens == max_new_tokens {
stop_reason = GenerationStopReason::TokenLimit;
break;
}
if guard.time_budget_exhausted(started.elapsed()) {
stop_reason = GenerationStopReason::TimeBudget;
break;
}
decode_tokens += 1;
let position = prompt_tokens + generated_tokens - 1;
let embedding = self.embedding(token)?;
let hidden = self.forward(
&embedding,
position,
cache,
scratch,
profile_enabled.then_some(&mut routes),
)?;
token = self.greedy_token(&hidden)?;
}
text.push_str(&decoder.finish()?);
let duration = started.elapsed();
let process = process_delta(profile_enabled, process_start);
Ok(DecodeOutput {
text,
generated_tokens,
decode_tokens,
#[cfg(feature = "bench")]
generated_token_ids,
duration,
stop_reason,
profile: PhaseProfile {
process,
expert_routes: routes,
},
})
}
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,
mut route_profile: Option<&mut ExpertRouteProfile>,
) -> 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, layer, route_profile.as_deref_mut())?;
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 prefill(
&self,
tokens: &[u32],
cache: &mut KvCache,
scratch: &mut AttentionScratch,
mut route_profile: Option<&mut ExpertRouteProfile>,
) -> Result<Vec<f32>> {
let ropes = (0..tokens.len())
.map(Rope::new)
.collect::<Result<Vec<_>>>()?;
let mut hidden = self.embeddings(tokens)?;
for (layer, weights) in self.layers.iter().enumerate() {
self.prefill_attention_layer(&mut hidden, &ropes, weights, layer, cache, scratch)?;
let feed_forward_norm = self.gguf.tensor_from_handle(&weights.feed_forward_norm);
let normalized = rms_norm(&hidden, HIDDEN, feed_forward_norm.f32_slice()?, EPSILON)?;
let mixture = self.mixture_of_experts_batch(
&normalized,
weights,
layer,
route_profile.as_deref_mut(),
)?;
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 prefill_attention_layer(
&self,
hidden: &mut [f32],
ropes: &[Rope],
weights: &LayerWeights,
layer: usize,
cache: &mut KvCache,
scratch: &mut AttentionScratch,
) -> Result<()> {
let expected_values = ropes
.len()
.checked_mul(HIDDEN)
.context("prefill hidden size overflow")?;
ensure!(
!ropes.is_empty() && hidden.len() == expected_values,
"prefill hidden shape differs"
);
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 batch_values = PREFILL_BATCH_SIZE * HIDDEN;
for (batch, batch_ropes) in hidden
.chunks_mut(batch_values)
.zip(ropes.chunks(PREFILL_BATCH_SIZE))
{
let row_count = batch_ropes.len();
let normalized = rms_norm(batch, 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(batch_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)?;
}
causal_gqa_batch(
&query,
&cache.layers[layer],
cached_tokens,
row_count,
sinks.f32_slice()?,
layer.is_multiple_of(2),
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(batch, &projected)?;
}
Ok(())
}
fn request_experts(
&self,
layer_index: usize,
layer: &LayerWeights,
experts: impl IntoIterator<Item = usize>,
) -> Result<Vec<ExpertRequest>> {
let requests = experts
.into_iter()
.map(|expert| {
let bundle = layer
.expert_bundles
.get(expert)
.with_context(|| format!("expert {expert} is out of range"))?;
Ok((ExpertKey::new(layer_index, expert), bundle.ranges))
})
.collect::<Result<Vec<_>>>()?;
self.expert_store.request_many(requests)
}
fn mixture_of_experts(
&self,
input: &[f32],
layer: &LayerWeights,
layer_index: usize,
route_profile: Option<&mut ExpertRouteProfile>,
) -> 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);
if let Some(profile) = route_profile {
profile.record(layer_index, &indices);
}
let expert_requests = self.request_experts(layer_index, layer, indices)?;
softmax(&mut weights);
let expert_outputs: Vec<Vec<f32>> = weights
.into_iter()
.zip(expert_requests)
.collect::<Vec<_>>()
.into_par_iter()
.map(|(routing_weight, request)| -> Result<Vec<f32>> {
let lease = request.wait()?;
let expert = layer.expert_weights(&lease)?;
let (mut gate, mut up) =
matrix_vector_pair(&expert.gate, &expert.up, input, Some(&input_activation))?;
add_bias(&mut gate, expert.gate_bias.f32_slice()?)?;
add_bias(&mut up, expert.up_bias.f32_slice()?)?;
swiglu_inplace(&mut gate, &up, SWIGLU_ALPHA, SWIGLU_LIMIT)?;
let mut output = matrix_vector(&expert.down, &gate)?;
add_bias(&mut output, expert.down_bias.f32_slice()?)?;
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 expert_batch_output(
input: &[f32],
routes: &[ExpertRoute],
request: ExpertRequest,
layer: &LayerWeights,
) -> Result<Vec<f32>> {
let lease = request.wait()?;
let expert = layer.expert_weights(&lease)?;
let output_values = routes
.len()
.checked_mul(HIDDEN)
.context("expert output size overflow")?;
let mut expert_output = Vec::with_capacity(output_values);
for route_batch in routes.chunks(PREFILL_BATCH_SIZE) {
let mut expert_input = Vec::with_capacity(route_batch.len() * HIDDEN);
for route in route_batch {
let start = route
.token
.checked_mul(HIDDEN)
.context("expert input offset overflow")?;
let end = start
.checked_add(HIDDEN)
.context("expert input range overflow")?;
expert_input.extend_from_slice(
input
.get(start..end)
.context("expert input route is out of range")?,
);
}
let (mut gate, mut up) =
matrix_matrix_pair(&expert.gate, &expert.up, &expert_input, route_batch.len())?;
add_bias_batch(&mut gate, expert.gate_bias.f32_slice()?)?;
add_bias_batch(&mut up, expert.up_bias.f32_slice()?)?;
swiglu_inplace(&mut gate, &up, SWIGLU_ALPHA, SWIGLU_LIMIT)?;
let mut output = matrix_matrix(&expert.down, &gate, route_batch.len())?;
add_bias_batch(&mut output, expert.down_bias.f32_slice()?)?;
for (values, route) in output.chunks_exact_mut(HIDDEN).zip(route_batch) {
for value in values {
*value *= route.weight;
}
}
expert_output.append(&mut output);
}
Ok(expert_output)
}
fn route_output_index(token: usize, slot: usize) -> Result<usize> {
ensure!(slot < ACTIVE_EXPERTS, "expert route slot is out of range");
token
.checked_mul(ACTIVE_EXPERTS)
.and_then(|index| index.checked_add(slot))
.context("expert route index overflow")
}
fn mix_expert_outputs(
input_len: usize,
row_count: usize,
routes_by_expert: &[Vec<ExpertRoute>],
expert_outputs: &[(usize, Vec<f32>)],
) -> Result<Vec<f32>> {
let route_count = row_count
.checked_mul(ACTIVE_EXPERTS)
.context("expert route count overflow")?;
let mut output_locations = vec![None; route_count];
for (output_index, (expert_index, output)) in expert_outputs.iter().enumerate() {
let expert_routes = routes_by_expert
.get(*expert_index)
.context("expert route batch is missing")?;
let expected_values = expert_routes
.len()
.checked_mul(HIDDEN)
.context("expert output shape overflow")?;
ensure!(
output.len() == expected_values,
"expert output shape differs"
);
for (expert_row, route) in expert_routes.iter().enumerate() {
let location_index = Self::route_output_index(route.token, route.slot)?;
let output_start = expert_row
.checked_mul(HIDDEN)
.context("expert output offset overflow")?;
let location = output_locations
.get_mut(location_index)
.context("expert output location is out of range")?;
ensure!(location.is_none(), "expert output location is duplicated");
*location = Some((output_index, output_start));
}
}
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 location_index = Self::route_output_index(token, slot)?;
let (output_index, start) = output_locations
.get(location_index)
.copied()
.flatten()
.context("routed expert output is missing")?;
let (_, output) = expert_outputs
.get(output_index)
.context("routed expert output is missing")?;
let end = start
.checked_add(HIDDEN)
.context("routed expert output range overflow")?;
let values = output
.get(start..end)
.context("routed expert output range is invalid")?;
for (value, routed) in mixed.iter_mut().zip(values) {
*value += routed;
}
}
}
Ok(mixture)
}
fn mixture_of_experts_batch(
&self,
input: &[f32],
layer: &LayerWeights,
layer_index: usize,
mut route_profile: Option<&mut ExpertRouteProfile>,
) -> 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);
if let Some(profile) = route_profile.as_deref_mut() {
profile.record(layer_index, &indices);
}
softmax(&mut weights);
for (slot, (expert, weight)) in indices.into_iter().zip(weights).enumerate() {
routes_by_expert[expert].push(ExpertRoute {
token,
slot,
weight,
});
}
}
let active_experts = routes_by_expert
.iter()
.enumerate()
.filter_map(|(expert, routes)| (!routes.is_empty()).then_some(expert))
.collect::<Vec<_>>();
let expert_requests =
self.request_experts(layer_index, layer, active_experts.iter().copied())?;
let expert_jobs = active_experts
.into_iter()
.zip(expert_requests)
.collect::<Vec<_>>();
let run_expert = |(expert_index, request)| -> Result<(usize, Vec<f32>)> {
let expert_routes = routes_by_expert
.get(expert_index)
.map(Vec::as_slice)
.context("expert route batch is missing")?;
let output = Self::expert_batch_output(input, expert_routes, request, layer)?;
Ok((expert_index, output))
};
let expert_outputs = if row_count < PREFILL_EXPERT_PARALLEL_ROWS {
expert_jobs
.into_iter()
.map(run_expert)
.collect::<Result<Vec<_>>>()?
} else {
expert_jobs
.into_par_iter()
.map(run_expert)
.collect::<Result<Vec<_>>>()?
};
Self::mix_expert_outputs(input.len(), row_count, &routes_by_expert, &expert_outputs)
}
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 head_dimension =
f32::from(u16::try_from(HEAD_DIMENSION).context("RoPE head dimension exceeds u16")?);
let low = correction_dimension(ROPE_BETA_FAST, head_dimension).max(0.0);
let high = correction_dimension(ROPE_BETA_SLOW, head_dimension).min(head_dimension - 1.0);
let mut rope = Self {
cosine: [0.0; HEAD_DIMENSION / 2],
sine: [0.0; HEAD_DIMENSION / 2],
};
for pair in 0..HEAD_DIMENSION / 2 {
let pair_f32 = f32::from(u16::try_from(pair).context("RoPE pair index exceeds u16")?);
let base_frequency = ROPE_BASE.powf(-(2.0 * pair_f32 / head_dimension));
let ramp = ((pair_f32 - 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, head_dimension: f32) -> 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 head_dimension =
f32::from(u16::try_from(HEAD_DIMENSION).context("attention head dimension exceeds u16")?);
let scale = 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()
.try_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 head_dimension =
f32::from(u16::try_from(HEAD_DIMENSION).context("attention head dimension exceeds u16")?);
let scale = 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()
.try_for_each(|(token, ((output, query), score_plane))| -> Result<()> {
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)?;
}
Ok(())
})?;
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],
) -> Result<()> {
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,
)?;
}
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(())
}