use std::path::Path;
use std::time::{Duration, Instant};
use anyhow::{Context as _, Result, ensure};
use rayon::prelude::*;
#[cfg(feature = "bench")]
use crate::alloc_profile::{AllocationProfile, AllocationSnapshot};
use crate::engine::model::gguf::{Gguf, Tensor, TensorHandle, TensorType};
use crate::engine::model::kernels::{
MatrixWorkspace, Q8Activation, dot_f32_fp16_group, matrix_argmax,
matrix_argmax_with_activation, matrix_matrix_into, matrix_matrix_pair_into,
matrix_matrix_triple_into, matrix_vector_q1_add_into, matrix_vector_q1_swiglu_quantize_into,
matrix_vector_q1_triple_into, rms_norm_into, rms_norm_quantize_into,
scale_accumulate_fp16_group, vector_add,
};
use crate::engine::model::profile::{ProcessDelta, ProcessSnapshot};
use crate::engine::model::tokenizer::{Tokenizer, Utf8Decoder};
use crate::engine::model::types::Fp16;
const LAYERS: usize = 36;
const HIDDEN: usize = 2_560;
const FEED_FORWARD: usize = 9_728;
const QUERY_HEADS: usize = 32;
const KEY_VALUE_HEADS: usize = 8;
const HEAD_DIMENSION: usize = 128;
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 VOCABULARY: usize = 151_669;
const PREFILL_BATCH_SIZE: usize = 128;
const EPSILON: f32 = 1.0e-6;
const ROPE_BASE: f32 = 5_000_000.0;
const ROPE_FACTOR: f32 = 4.0;
const ROPE_ORIGINAL_CONTEXT: usize = 8_192;
const ROPE_ORIGINAL_CONTEXT_F32: f32 = 8_192.0;
const ROPE_BETA_FAST: f32 = 32.0;
const ROPE_BETA_SLOW: f32 = 1.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 Q1_BLOCK_VALUES: usize = 128;
const Q1_BLOCK_BYTES: usize = 18;
const Q1_SIGN_BYTES: usize = 16;
#[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>,
#[cfg(feature = "bench")]
pub prefill_allocations: AllocationProfile,
#[cfg(feature = "bench")]
pub decode_allocations: AllocationProfile,
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>,
token_embd: TensorHandle,
output_norm: TensorHandle,
output: TensorHandle,
}
struct LayerWeights {
attention_norm: TensorHandle,
query: TensorHandle,
query_norm: TensorHandle,
key: TensorHandle,
key_norm: TensorHandle,
value: TensorHandle,
attention_output: TensorHandle,
feed_forward_norm: TensorHandle,
gate: TensorHandle,
up: TensorHandle,
down: TensorHandle,
}
struct LayerNames {
attention_norm: String,
query: String,
query_norm: String,
key: String,
key_norm: String,
value: String,
attention_output: String,
feed_forward_norm: String,
gate: String,
up: String,
down: String,
}
pub(super) struct PhaseProfile {
pub(super) process: Option<ProcessDelta>,
}
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 {
prompt_tokens: usize,
max_new_tokens: usize,
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_norm: format!("{prefix}.attn_q_norm.weight"),
key: format!("{prefix}.attn_k.weight"),
key_norm: format!("{prefix}.attn_k_norm.weight"),
value: format!("{prefix}.attn_v.weight"),
attention_output: format!("{prefix}.attn_output.weight"),
feed_forward_norm: format!("{prefix}.ffn_norm.weight"),
gate: format!("{prefix}.ffn_gate.weight"),
up: format!("{prefix}.ffn_up.weight"),
down: format!("{prefix}.ffn_down.weight"),
}
}
}
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_norm: gguf.resolve(&names.query_norm)?,
key: gguf.resolve(&names.key)?,
key_norm: gguf.resolve(&names.key_norm)?,
value: gguf.resolve(&names.value)?,
attention_output: gguf.resolve(&names.attention_output)?,
feed_forward_norm: gguf.resolve(&names.feed_forward_norm)?,
gate: gguf.resolve(&names.gate)?,
up: gguf.resolve(&names.up)?,
down: gguf.resolve(&names.down)?,
})
}
}
impl TextModel {
pub fn load(path: impl AsRef<Path>) -> Result<Self> {
let path = path.as_ref();
let gguf = Gguf::load(path).context("load Qwen3 GGUF")?;
validate_metadata(&gguf).context("validate Qwen3 metadata")?;
validate_tensors(&gguf).context("validate Qwen3 tensors")?;
let tokenizer = Tokenizer::from_gguf(&gguf).context("load Qwen3 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 = token_embd.clone();
Ok(Self {
gguf,
tokenizer,
layers,
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();
#[cfg(feature = "bench")]
let prefill_allocations = AllocationSnapshot::now();
let mut cache = KvCache::new(max_context_tokens)?;
let mut attention_scratch = AttentionScratch::new(max_context_tokens)?;
let mut prefill_scratch = PrefillScratch::new(prompt_tokens)?;
let prefill_process_start = process_snapshot(profile_enabled);
let prefill_started = Instant::now();
self.prefill(
&prompt_ids,
&mut cache,
&mut attention_scratch,
&mut prefill_scratch,
)?;
if max_new_tokens == 0 {
let prefill_duration = prefill_started.elapsed();
let prefill_process = process_delta(profile_enabled, prefill_process_start);
#[cfg(feature = "bench")]
let prefill_allocations = prefill_allocations.elapsed();
return Ok(Self::empty_generation(
prompt_tokens,
prefill_duration,
prefill_process,
profile_enabled,
#[cfg(feature = "bench")]
prefill_allocations,
));
}
let last_start = prefill_scratch
.normalized
.len()
.checked_sub(HIDDEN)
.context("prefill output is shorter than the hidden width")?;
let last_hidden = prefill_scratch
.normalized
.get(last_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);
#[cfg(feature = "bench")]
let prefill_allocations = prefill_allocations.elapsed();
let mut decode_scratch = DecodeScratch::new()?;
#[cfg(feature = "bench")]
let decode_allocations = AllocationSnapshot::now();
let decoded = self.decode(
token,
&mut cache,
&mut attention_scratch,
&mut decode_scratch,
DecodeConfig {
prompt_tokens,
max_new_tokens,
profile_enabled,
max_duration: max_decode_duration,
},
)?;
#[cfg(feature = "bench")]
let decode_allocations = decode_allocations.elapsed();
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,
#[cfg(feature = "bench")]
prefill_allocations,
#[cfg(feature = "bench")]
decode_allocations,
prefill_duration,
decode_duration: decoded.duration,
stop_reason: decoded.stop_reason,
profile: profile_enabled.then_some(GenerationProfile {
prefill: PhaseProfile {
process: prefill_process,
},
decode: decoded.profile,
}),
})
}
fn empty_generation(
prompt_tokens: usize,
prefill_duration: Duration,
prefill_process: Option<ProcessDelta>,
profile_enabled: bool,
#[cfg(feature = "bench")] prefill_allocations: AllocationProfile,
) -> Generation {
Generation {
text: String::new(),
prompt_tokens,
generated_tokens: 0,
decode_tokens: 0,
#[cfg(feature = "bench")]
generated_token_ids: Vec::new(),
#[cfg(feature = "bench")]
prefill_allocations,
#[cfg(feature = "bench")]
decode_allocations: AllocationProfile::default(),
prefill_duration,
decode_duration: Duration::ZERO,
stop_reason: GenerationStopReason::TokenLimit,
profile: profile_enabled.then_some(GenerationProfile {
prefill: PhaseProfile {
process: prefill_process,
},
decode: PhaseProfile { process: None },
}),
}
}
fn decode(
&self,
mut token: u32,
cache: &mut KvCache,
attention_scratch: &mut AttentionScratch,
scratch: &mut DecodeScratch,
config: DecodeConfig,
) -> Result<DecodeOutput> {
let DecodeConfig {
prompt_tokens,
max_new_tokens,
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 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;
self.embedding_into(token, &mut scratch.hidden)?;
self.forward(position, cache, attention_scratch, scratch)?;
token = self.greedy_token_with_activation(&scratch.hidden, &scratch.activation)?;
}
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 },
})
}
fn embedding_into(&self, token: u32, output: &mut [f32]) -> Result<()> {
ensure!(output.len() == HIDDEN, "embedding output width differs");
let tensor = self.gguf.tensor_from_handle(&self.token_embd);
let row = usize::try_from(token).context("token ID exceeds usize")?;
dequantize_q1_row(tensor.encoded_row(row)?, output).context("dequantize embedding row")
}
fn embeddings_into(&self, tokens: &[u32], output: &mut [f32]) -> Result<()> {
let output_len = tokens
.len()
.checked_mul(HIDDEN)
.context("embedding batch size overflow")?;
ensure!(
output.len() == output_len,
"embedding batch output shape differs"
);
let tensor = self.gguf.tensor_from_handle(&self.token_embd);
for (token, embedding) in tokens.iter().zip(output.chunks_exact_mut(HIDDEN)) {
let row = usize::try_from(*token).context("token ID exceeds usize")?;
dequantize_q1_row(tensor.encoded_row(row)?, embedding)
.with_context(|| format!("dequantize embedding row {row}"))?;
}
Ok(())
}
fn forward(
&self,
position: usize,
cache: &mut KvCache,
attention_scratch: &mut AttentionScratch,
scratch: &mut DecodeScratch,
) -> Result<()> {
let rope = Rope::new(position)?;
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_norm = self.gguf.tensor_from_handle(&weights.query_norm);
let key_norm = self.gguf.tensor_from_handle(&weights.key_norm);
let attention_output = self.gguf.tensor_from_handle(&weights.attention_output);
let feed_forward_norm = self.gguf.tensor_from_handle(&weights.feed_forward_norm);
let gate_w = self.gguf.tensor_from_handle(&weights.gate);
let up_w = self.gguf.tensor_from_handle(&weights.up);
let down_w = self.gguf.tensor_from_handle(&weights.down);
{
let DecodeScratch {
hidden,
query,
key,
value,
activation,
} = scratch;
rms_norm_quantize_into(
hidden,
HIDDEN,
attention_norm.f32_slice()?,
EPSILON,
activation,
)?;
matrix_vector_q1_triple_into(
[&query_w, &key_w, &value_w],
hidden,
activation,
[query, key, value],
)?;
apply_head_rms_norm_rope(query, QUERY_HEADS, query_norm.f32_slice()?, &rope)?;
apply_head_rms_norm_rope(key, KEY_VALUE_HEADS, key_norm.f32_slice()?, &rope)?;
cache.layers[layer].append(key, value)?;
causal_gqa(
query,
&cache.layers[layer],
cache.layers[layer].token_count(),
attention_scratch,
activation,
)?;
matrix_vector_q1_add_into(&attention_output, activation, hidden)?;
rms_norm_quantize_into(
hidden,
HIDDEN,
feed_forward_norm.f32_slice()?,
EPSILON,
activation,
)?;
matrix_vector_q1_swiglu_quantize_into(&gate_w, &up_w, activation)?;
matrix_vector_q1_add_into(&down_w, activation, hidden)?;
}
}
let output_norm = self.gguf.tensor_from_handle(&self.output_norm);
rms_norm_quantize_into(
&scratch.hidden,
HIDDEN,
output_norm.f32_slice()?,
EPSILON,
&mut scratch.activation,
)
}
fn prefill(
&self,
tokens: &[u32],
cache: &mut KvCache,
attention_scratch: &mut AttentionScratch,
scratch: &mut PrefillScratch,
) -> Result<()> {
let ropes = (0..tokens.len())
.map(Rope::new)
.collect::<Result<Vec<_>>>()?;
self.embeddings_into(tokens, &mut scratch.hidden)?;
for (layer, weights) in self.layers.iter().enumerate() {
self.prefill_attention_layer(
&ropes,
weights,
layer,
cache,
attention_scratch,
scratch,
)?;
let feed_forward_norm = self.gguf.tensor_from_handle(&weights.feed_forward_norm);
let gate_w = self.gguf.tensor_from_handle(&weights.gate);
let up_w = self.gguf.tensor_from_handle(&weights.up);
let down_w = self.gguf.tensor_from_handle(&weights.down);
let token_count = tokens.len();
let batch_values = PREFILL_BATCH_SIZE
.checked_mul(HIDDEN)
.context("prefill feed-forward batch overflow")?;
for batch in scratch.hidden.chunks_mut(batch_values) {
let row_count = batch.len() / HIDDEN;
let hidden_values = row_count
.checked_mul(HIDDEN)
.context("prefill feed-forward hidden overflow")?;
let ff_values = row_count
.checked_mul(FEED_FORWARD)
.context("prefill feed-forward width overflow")?;
let normalized = &mut scratch.normalized[..hidden_values];
let gate = &mut scratch.gate[..ff_values];
let up = &mut scratch.up[..ff_values];
let projected = &mut scratch.projected[..hidden_values];
rms_norm_into(
batch,
HIDDEN,
feed_forward_norm.f32_slice()?,
EPSILON,
normalized,
)?;
matrix_matrix_pair_into(
&gate_w,
&up_w,
normalized,
row_count,
gate,
up,
&mut scratch.feed_forward_workspace,
)?;
for (gate_row, up_row) in gate
.chunks_exact_mut(FEED_FORWARD)
.zip(up.chunks_exact(FEED_FORWARD))
{
dense_swiglu_inplace(gate_row, up_row)?;
}
matrix_matrix_into(
&down_w,
gate,
row_count,
projected,
&mut scratch.feed_forward_output_workspace,
)?;
vector_add(batch, projected)?;
}
ensure!(
scratch.hidden.len()
== token_count
.checked_mul(HIDDEN)
.context("prefill hidden length overflow")?,
"prefill hidden shape differs after feed-forward"
);
}
let output_norm = self.gguf.tensor_from_handle(&self.output_norm);
rms_norm_into(
&scratch.hidden,
HIDDEN,
output_norm.f32_slice()?,
EPSILON,
&mut scratch.normalized,
)
}
fn prefill_attention_layer(
&self,
ropes: &[Rope],
weights: &LayerWeights,
layer: usize,
cache: &mut KvCache,
attention_scratch: &mut AttentionScratch,
scratch: &mut PrefillScratch,
) -> Result<()> {
let expected_values = ropes
.len()
.checked_mul(HIDDEN)
.context("prefill hidden size overflow")?;
ensure!(
!ropes.is_empty() && scratch.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_norm = self.gguf.tensor_from_handle(&weights.query_norm);
let key_norm = self.gguf.tensor_from_handle(&weights.key_norm);
let attention_output = self.gguf.tensor_from_handle(&weights.attention_output);
let batch_values = PREFILL_BATCH_SIZE * HIDDEN;
let PrefillScratch {
hidden,
attention_normalized,
query,
key,
value,
projected,
attention_qkv_workspace,
attention_output_workspace,
..
} = scratch;
for (batch, batch_ropes) in hidden
.chunks_mut(batch_values)
.zip(ropes.chunks(PREFILL_BATCH_SIZE))
{
let row_count = batch_ropes.len();
let hidden_values = row_count * HIDDEN;
let query_values = row_count * QUERY_SIZE;
let key_value_values = row_count * KEY_VALUE_SIZE;
let normalized = &mut attention_normalized[..hidden_values];
let query = &mut query[..query_values];
let key = &mut key[..key_value_values];
let value = &mut value[..key_value_values];
let projected = &mut projected[..hidden_values];
rms_norm_into(
batch,
HIDDEN,
attention_norm.f32_slice()?,
EPSILON,
normalized,
)?;
matrix_matrix_triple_into(
[&query_w, &key_w, &value_w],
normalized,
row_count,
[query, key, value],
attention_qkv_workspace,
)?;
for ((query_row, key_row), rope) in query
.chunks_exact_mut(QUERY_SIZE)
.zip(key.chunks_exact_mut(KEY_VALUE_SIZE))
.zip(batch_ropes)
{
apply_head_rms_norm_rope(query_row, QUERY_HEADS, query_norm.f32_slice()?, rope)?;
apply_head_rms_norm_rope(key_row, KEY_VALUE_HEADS, key_norm.f32_slice()?, rope)?;
}
let cached_tokens = cache.layers[layer].token_count();
for (key_row, value_row) in key
.chunks_exact(KEY_VALUE_SIZE)
.zip(value.chunks_exact(KEY_VALUE_SIZE))
{
cache.layers[layer].append(key_row, value_row)?;
}
causal_gqa_batch(
query,
&cache.layers[layer],
cached_tokens,
row_count,
attention_scratch,
)?;
matrix_matrix_into(
&attention_output,
attention_scratch.output_rows(row_count),
row_count,
projected,
attention_output_workspace,
)?;
vector_add(batch, projected)?;
}
Ok(())
}
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 greedy_token_with_activation(
&self,
hidden: &[f32],
activation: &Q8Activation,
) -> Result<u32> {
let output = self.gguf.tensor_from_handle(&self.output);
let token = matrix_argmax_with_activation(&output, hidden, activation)
.context("compute greedy token")?;
u32::try_from(token).context("token ID exceeds u32")
}
}
fn dequantize_q1_row(row: &[u8], output: &mut [f32]) -> Result<()> {
ensure!(
output.len().is_multiple_of(Q1_BLOCK_VALUES),
"Q1_0 output width is not divisible by 128"
);
let groups = output.len() / Q1_BLOCK_VALUES;
ensure!(
row.len()
== groups
.checked_mul(Q1_BLOCK_BYTES)
.context("Q1_0 row size overflow")?,
"invalid Q1_0 row size"
);
for (block, values) in row
.chunks_exact(Q1_BLOCK_BYTES)
.zip(output.chunks_exact_mut(Q1_BLOCK_VALUES))
{
let scale = Fp16::decode_le(&block[..2]).to_f32();
let signs = &block[2..];
debug_assert_eq!(signs.len(), Q1_SIGN_BYTES);
for (byte_index, byte) in signs.iter().copied().enumerate() {
let base = byte_index * 8;
for bit in 0..8 {
values[base + bit] = if byte & (1u8 << bit) != 0 {
scale
} else {
-scale
};
}
}
}
Ok(())
}
fn dense_swiglu_inplace(gate: &mut [f32], up: &[f32]) -> Result<()> {
ensure!(gate.len() == up.len(), "swiglu widths differ");
for (gate, up) in gate.iter_mut().zip(up) {
let silu = *gate / (1.0 + (-*gate).exp());
*gate = silu * *up;
}
Ok(())
}
fn apply_head_rms_norm_rope(
values: &mut [f32],
head_count: usize,
weight: &[f32],
rope: &Rope,
) -> Result<()> {
ensure!(
weight.len() == HEAD_DIMENSION,
"head RMS norm weight width differs"
);
let expected = head_count
.checked_mul(HEAD_DIMENSION)
.context("head RMS norm size overflow")?;
ensure!(
values.len() == expected,
"head RMS norm input width differs"
);
let half = HEAD_DIMENSION / 2;
let (weight_first, weight_second) = weight.split_at(half);
for head in values.chunks_exact_mut(HEAD_DIMENSION) {
let mut mean_square = 0.0f32;
for value in head.iter() {
mean_square += value * value;
}
mean_square /=
f32::from(u16::try_from(HEAD_DIMENSION).context("head dimension exceeds u16")?);
let scale = (mean_square + EPSILON).sqrt().recip();
let (first, second) = head.split_at_mut(half);
for (((((first, second), cosine), sine), weight_first), weight_second) in first
.iter_mut()
.zip(second.iter_mut())
.zip(&rope.cosine)
.zip(&rope.sine)
.zip(weight_first)
.zip(weight_second)
{
let left = *first * scale * *weight_first;
let right = *second * scale * *weight_second;
*first = left * cosine - right * sine;
*second = left * sine + right * cosine;
}
}
Ok(())
}
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 u16")?,
))
}
fn correction_dimension(rotations: f32, head_dimension: f32) -> f32 {
head_dimension * (ROPE_ORIGINAL_CONTEXT_F32 / (rotations * 2.0 * std::f32::consts::PI)).ln()
/ (2.0 * ROPE_BASE.ln())
}
#[derive(Default)]
struct LayerCache {
keys: [Vec<Fp16>; KEY_VALUE_HEADS],
values: [Vec<Fp16>; KEY_VALUE_HEADS],
}
impl LayerCache {
fn reserve(&mut self, tokens: usize) -> Result<()> {
let values = tokens
.checked_mul(HEAD_DIMENSION)
.context("key/value cache size overflow")?;
for (keys, values_head) in self.keys.iter_mut().zip(&mut self.values) {
keys.try_reserve_exact(values)
.context("reserve key cache")?;
values_head
.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"
);
for (keys, key_head) in self.keys.iter_mut().zip(key.chunks_exact(HEAD_DIMENSION)) {
keys.extend(key_head.iter().copied().map(Fp16::from));
}
for (values, value_head) in self
.values
.iter_mut()
.zip(value.chunks_exact(HEAD_DIMENSION))
{
values.extend(value_head.iter().copied().map(Fp16::from));
}
Ok(())
}
fn token_count(&self) -> usize {
self.keys
.first()
.map_or(0, |keys| keys.len() / HEAD_DIMENSION)
}
fn consistent(&self) -> bool {
let expected = self.keys.first().map_or(0, Vec::len);
self.keys
.iter()
.chain(&self.values)
.all(|head| head.len() == expected)
}
fn head_range(
&self,
key_value_head: usize,
first_token: usize,
tokens: usize,
) -> Result<(&[Fp16], &[Fp16])> {
let start = first_token
.checked_mul(HEAD_DIMENSION)
.context("attention cache index overflow")?;
let values = tokens
.checked_mul(HEAD_DIMENSION)
.context("attention cache range overflow")?;
let end = start
.checked_add(values)
.context("attention cache range overflow")?;
let key = self
.keys
.get(key_value_head)
.and_then(|keys| keys.get(start..end))
.context("attention key cache slice out of range")?;
let value = self
.values
.get(key_value_head)
.and_then(|values| values.get(start..end))
.context("attention value cache slice out of range")?;
Ok((key, value))
}
}
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>,
inverses: [f32; QUERY_HEADS],
}
impl AttentionScratch {
fn new(tokens: usize) -> Result<Self> {
ensure!(tokens != 0, "attention cache capacity must not be zero");
let output_capacity = PREFILL_BATCH_SIZE
.checked_mul(QUERY_SIZE)
.context("attention output workspace overflow")?;
let mut output = Vec::new();
output
.try_reserve_exact(output_capacity)
.context("reserve attention output")?;
output.resize(output_capacity, 0.0);
Ok(Self {
output,
inverses: [0.0; QUERY_HEADS],
})
}
fn output_rows(&self, row_count: usize) -> &[f32] {
&self.output[..row_count * QUERY_SIZE]
}
fn batch_output(&mut self, row_count: usize) -> Result<&mut [f32]> {
let output_values = row_count
.checked_mul(QUERY_SIZE)
.context("attention batch output overflow")?;
ensure!(
output_values <= self.output.len(),
"attention batch exceeds output workspace"
);
Ok(&mut self.output[..output_values])
}
}
struct PrefillScratch {
hidden: Vec<f32>,
normalized: Vec<f32>,
attention_normalized: Vec<f32>,
query: Vec<f32>,
key: Vec<f32>,
value: Vec<f32>,
projected: Vec<f32>,
gate: Vec<f32>,
up: Vec<f32>,
attention_qkv_workspace: MatrixWorkspace,
attention_output_workspace: MatrixWorkspace,
feed_forward_workspace: MatrixWorkspace,
feed_forward_output_workspace: MatrixWorkspace,
}
impl PrefillScratch {
fn new(tokens: usize) -> Result<Self> {
let hidden_values = tokens
.checked_mul(HIDDEN)
.context("prefill hidden workspace overflow")?;
let batch_rows = tokens.min(PREFILL_BATCH_SIZE);
let batch_hidden = batch_rows
.checked_mul(HIDDEN)
.context("prefill batch hidden workspace overflow")?;
let batch_query = batch_rows
.checked_mul(QUERY_SIZE)
.context("prefill batch query workspace overflow")?;
let batch_key_value = batch_rows
.checked_mul(KEY_VALUE_SIZE)
.context("prefill batch key/value workspace overflow")?;
let batch_feed_forward = batch_rows
.checked_mul(FEED_FORWARD)
.context("prefill batch feed-forward workspace overflow")?;
Ok(Self {
hidden: vec![0.0; hidden_values],
normalized: vec![0.0; hidden_values],
attention_normalized: vec![0.0; batch_hidden],
query: vec![0.0; batch_query],
key: vec![0.0; batch_key_value],
value: vec![0.0; batch_key_value],
projected: vec![0.0; batch_hidden],
gate: vec![0.0; batch_feed_forward],
up: vec![0.0; batch_feed_forward],
attention_qkv_workspace: MatrixWorkspace::default(),
attention_output_workspace: MatrixWorkspace::default(),
feed_forward_workspace: MatrixWorkspace::default(),
feed_forward_output_workspace: MatrixWorkspace::default(),
})
}
}
struct DecodeScratch {
hidden: Vec<f32>,
query: Vec<f32>,
key: Vec<f32>,
value: Vec<f32>,
activation: Q8Activation,
}
impl DecodeScratch {
fn new() -> Result<Self> {
Ok(Self {
hidden: vec![0.0; HIDDEN],
query: vec![0.0; QUERY_SIZE],
key: vec![0.0; KEY_VALUE_SIZE],
value: vec![0.0; KEY_VALUE_SIZE],
activation: Q8Activation::with_capacity(FEED_FORWARD)?,
})
}
}
fn causal_gqa(
query: &[f32],
cache: &LayerCache,
token_count: usize,
scratch: &mut AttentionScratch,
activation: &mut Q8Activation,
) -> Result<()> {
ensure!(query.len() == QUERY_SIZE, "attention query width differs");
ensure!(cache.consistent(), "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 output = &mut scratch.output[..QUERY_SIZE];
let inverses = &mut scratch.inverses;
let head_dimension =
f32::from(u16::try_from(HEAD_DIMENSION).context("attention head dimension exceeds u16")?);
let scale = head_dimension.sqrt().recip();
let ctx = AttentionHeadCtx {
cache,
first_token: 0,
visible_tokens: token_count,
scale,
};
inverses
.par_chunks_exact_mut(QUERY_GROUP)
.zip(output.par_chunks_exact_mut(HEAD_DIMENSION * QUERY_GROUP))
.zip(query.par_chunks_exact(HEAD_DIMENSION * QUERY_GROUP))
.enumerate()
.try_for_each(
|(key_value_head, ((inverses, output), query_values))| -> Result<()> {
inverses.copy_from_slice(&attention_group(
key_value_head,
query_values,
&ctx,
output,
false,
)?);
Ok(())
},
)?;
ensure!(
output.iter().all(|value| value.is_finite()),
"attention output is not finite"
);
activation.quantize_into(output)?;
activation.scale_repeating(HEAD_DIMENSION, inverses)
}
fn causal_gqa_batch(
queries: &[f32],
cache: &LayerCache,
cached_tokens: usize,
row_count: usize,
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!(cache.consistent(), "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 = scratch.batch_output(row_count)?;
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))
.enumerate()
.try_for_each(|(token, (output, query))| -> Result<()> {
let token_count = cached_tokens + token + 1;
let ctx = AttentionHeadCtx {
cache,
first_token: 0,
visible_tokens: token_count,
scale,
};
for key_value_head in 0..KEY_VALUE_HEADS {
let query_start = key_value_head
.checked_mul(QUERY_GROUP)
.and_then(|base| base.checked_mul(HEAD_DIMENSION))
.context("attention group offset overflow")?;
let query_end = query_start
.checked_add(QUERY_GROUP * HEAD_DIMENSION)
.context("attention group width overflow")?;
let query_values = query
.get(query_start..query_end)
.context("attention group query slice out of range")?;
let head_output = output
.get_mut(query_start..query_end)
.context("attention group output slice out of range")?;
attention_group(key_value_head, query_values, &ctx, head_output, true)?;
}
Ok(())
})?;
ensure!(
outputs.iter().all(|value| value.is_finite()),
"attention output is not finite"
);
Ok(())
}
struct AttentionHeadCtx<'a> {
cache: &'a LayerCache,
first_token: usize,
visible_tokens: usize,
scale: f32,
}
fn attention_group(
key_value_head: usize,
query_values: &[f32],
ctx: &AttentionHeadCtx<'_>,
output: &mut [f32],
apply_inverse: bool,
) -> Result<[f32; QUERY_GROUP]> {
let group_values = QUERY_GROUP
.checked_mul(HEAD_DIMENSION)
.context("attention group width overflow")?;
ensure!(
query_values.len() == group_values,
"attention group query width differs"
);
ensure!(
output.len() == group_values,
"attention group output width differs"
);
let (keys, values) =
ctx.cache
.head_range(key_value_head, ctx.first_token, ctx.visible_tokens)?;
let mut maximums = [f32::NEG_INFINITY; QUERY_GROUP];
let mut sums = [0.0_f32; QUERY_GROUP];
output.fill(0.0);
for (key, value) in keys
.chunks_exact(HEAD_DIMENSION)
.zip(values.chunks_exact(HEAD_DIMENSION))
{
let scores = dot_f32_fp16_group(query_values, key)?;
let mut rescales = [0.0_f32; QUERY_GROUP];
let mut weights = [0.0_f32; QUERY_GROUP];
for (((&score, maximum), sum), (rescale, weight)) in scores
.iter()
.zip(&mut maximums)
.zip(&mut sums)
.zip(rescales.iter_mut().zip(weights.iter_mut()))
{
let score = score * ctx.scale;
let next_maximum = (*maximum).max(score);
*rescale = (*maximum - next_maximum).exp();
*weight = (score - next_maximum).exp();
*sum = *sum * *rescale + *weight;
*maximum = next_maximum;
}
scale_accumulate_fp16_group(output, value, rescales, weights)?;
}
let mut inverses = [0.0_f32; QUERY_GROUP];
for ((dest, &sum), inverse) in output
.chunks_exact_mut(HEAD_DIMENSION)
.zip(&sums)
.zip(&mut inverses)
{
*inverse = if sum == 0.0 { 0.0 } else { sum.recip() };
if apply_inverse {
for value in dest {
*value *= *inverse;
}
}
}
Ok(inverses)
}
fn validate_metadata(gguf: &Gguf) -> Result<()> {
ensure!(
gguf.architecture()? == "qwen3",
"expected qwen3 architecture"
);
validate_u32(gguf, "qwen3.block_count", LAYERS)?;
validate_u32(gguf, "qwen3.embedding_length", HIDDEN)?;
validate_u32(gguf, "qwen3.feed_forward_length", FEED_FORWARD)?;
validate_u32(gguf, "qwen3.attention.head_count", QUERY_HEADS)?;
validate_u32(gguf, "qwen3.attention.head_count_kv", KEY_VALUE_HEADS)?;
validate_u32(gguf, "qwen3.attention.key_length", HEAD_DIMENSION)?;
validate_u32(gguf, "qwen3.attention.value_length", HEAD_DIMENSION)?;
let epsilon = gguf.f32("qwen3.attention.layer_norm_rms_epsilon")?;
ensure!(
(epsilon - EPSILON).abs() <= 1.0e-12,
"RMS epsilon is {epsilon}, expected {EPSILON}"
);
let rope_base = gguf.f32("qwen3.rope.freq_base")?;
ensure!(
rope_base.to_bits() == ROPE_BASE.to_bits(),
"RoPE base is {rope_base}, expected {ROPE_BASE}"
);
let rope_factor = gguf.f32("qwen3.rope.scaling.factor")?;
ensure!(
rope_factor.to_bits() == ROPE_FACTOR.to_bits(),
"RoPE scaling factor is {rope_factor}, expected {ROPE_FACTOR}"
);
validate_u32(
gguf,
"qwen3.rope.scaling.original_context_length",
ROPE_ORIGINAL_CONTEXT,
)?;
ensure!(
gguf.string("qwen3.rope.scaling.type")? == "yarn",
"expected YaRN rope scaling"
);
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")? == "qwen2",
"expected qwen2 pre-tokenizer"
);
ensure!(
!gguf.bool("tokenizer.ggml.add_bos_token")?,
"Qwen3 must not prepend BOS"
);
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::Q1_0,
)?;
validate_kind(gguf, "output_norm.weight", &[HIDDEN], TensorType::F32)?;
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::Q1_0)?;
validate_kind(gguf, &names.query_norm, &[HEAD_DIMENSION], TensorType::F32)?;
validate_kind(
gguf,
&names.key,
&[HIDDEN, KEY_VALUE_SIZE],
TensorType::Q1_0,
)?;
validate_kind(gguf, &names.key_norm, &[HEAD_DIMENSION], TensorType::F32)?;
validate_kind(
gguf,
&names.value,
&[HIDDEN, KEY_VALUE_SIZE],
TensorType::Q1_0,
)?;
validate_kind(
gguf,
&names.attention_output,
&[QUERY_SIZE, HIDDEN],
TensorType::Q1_0,
)?;
validate_kind(gguf, &names.feed_forward_norm, &[HIDDEN], TensorType::F32)?;
validate_kind(gguf, &names.gate, &[HIDDEN, FEED_FORWARD], TensorType::Q1_0)?;
validate_kind(gguf, &names.up, &[HIDDEN, FEED_FORWARD], TensorType::Q1_0)?;
validate_kind(gguf, &names.down, &[FEED_FORWARD, HIDDEN], TensorType::Q1_0)?;
}
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()
);
let _ = Tensor::tensor_type(&tensor);
Ok(())
}