diff --git a/CHANGELOG.md b/CHANGELOG.md index d053b68..24ba968 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,9 @@ versioning for published crates while it remains in the `0.x` series. test case. - Reused the existing token buffer allocation when applying BPE merges, while preserving the exact tokenizer output contract. +- Reused one normalization scratch buffer across every transformer layer in a + single-stream token step, eliminating repeated layer-local allocations while + preserving the exact generated token trace. - Applied the canonical Rust 2024 rustfmt style across the workspace. ### Documentation diff --git a/crates/ferrite-inference/src/scalar/math.rs b/crates/ferrite-inference/src/scalar/math.rs index 81eac44..9f49bb2 100644 --- a/crates/ferrite-inference/src/scalar/math.rs +++ b/crates/ferrite-inference/src/scalar/math.rs @@ -7,6 +7,17 @@ use super::InferenceError; /// Returns an error for empty or mismatched inputs, non-finite values, an /// invalid epsilon, or a zero or non-finite normalization scale. pub fn rms_norm(input: &[f32], weight: &[f32], epsilon: f32) -> Result, InferenceError> { + let mut output = Vec::with_capacity(input.len()); + rms_norm_into(input, weight, epsilon, &mut output)?; + Ok(output) +} + +pub(super) fn rms_norm_into( + input: &[f32], + weight: &[f32], + epsilon: f32, + output: &mut Vec, +) -> Result<(), InferenceError> { if input.is_empty() { return Err(InferenceError::new("rms_norm input must not be empty")); } @@ -32,7 +43,8 @@ pub fn rms_norm(input: &[f32], weight: &[f32], epsilon: f32) -> Result, return Err(InferenceError::new("rms_norm scale must be finite")); } - let mut output = Vec::with_capacity(input.len()); + output.clear(); + output.reserve(input.len()); for (value, weight) in input.iter().zip(weight.iter()) { let normalized = value / scale * weight; if !normalized.is_finite() { @@ -40,7 +52,7 @@ pub fn rms_norm(input: &[f32], weight: &[f32], epsilon: f32) -> Result, } output.push(normalized); } - Ok(output) + Ok(()) } /// Returns the index of the first greatest finite value. @@ -209,6 +221,21 @@ mod tests { Ok(()) } + #[test] + fn rms_norm_into_reuses_output_allocation() -> Result<(), InferenceError> { + let mut output = Vec::new(); + rms_norm_into(&[3.0, 4.0], &[1.0, 0.5], 0.0, &mut output)?; + let allocation = output.as_ptr(); + let capacity = output.capacity(); + + rms_norm_into(&[1.0, -1.0], &[0.5, 0.25], 0.0, &mut output)?; + + assert_eq!(output.as_ptr(), allocation); + assert_eq!(output.capacity(), capacity); + assert_eq!(output, [0.5, -0.25]); + Ok(()) + } + #[test] fn softmax_in_place_preserves_reference_operation_order() -> Result<(), InferenceError> { let mut values = vec![-3.0f32, 0.25, 1.5, -0.75]; diff --git a/crates/ferrite-inference/src/scalar/session.rs b/crates/ferrite-inference/src/scalar/session.rs index d9abb61..d3c5ae1 100644 --- a/crates/ferrite-inference/src/scalar/session.rs +++ b/crates/ferrite-inference/src/scalar/session.rs @@ -1,7 +1,7 @@ use super::{ InferenceError, NextToken, Q8KActivationMatvecRole, ScalarExecutionOptions, ScalarLlamaModel, attention::causal_attention, - math::{add_assign, argmax, rms_norm, swiglu_in_place}, + math::{add_assign, argmax, rms_norm, rms_norm_into, swiglu_in_place}, profile::{ProfiledNextToken, ProfiledTokenId, ScalarMatVecComparison, ScalarProfileEvent}, }; @@ -384,15 +384,17 @@ impl<'a> ScalarLlamaSession<'a> { let position = self.cached_token_count; let mut hidden = self.model.weights.token_embedding.row_values(token_id)?; + let mut normed = Vec::with_capacity(self.model.config.hidden_size); for (layer_index, layer) in self.model.weights.layers.iter().enumerate() { if on_layer(layer_index)? == PromptEvaluationControl::Cancel { return Err(InferenceError::new("prompt evaluation cancelled")); } - let normed = rms_norm( + rms_norm_into( &hidden, &layer.attn_norm, self.model.config.rms_norm_epsilon, + &mut normed, )?; let q_options = self .options @@ -487,8 +489,12 @@ impl<'a> ScalarLlamaSession<'a> { )?; add_assign(&mut hidden, &attention_output)?; - let ffn_normed = - rms_norm(&hidden, &layer.ffn_norm, self.model.config.rms_norm_epsilon)?; + rms_norm_into( + &hidden, + &layer.ffn_norm, + self.model.config.rms_norm_epsilon, + &mut normed, + )?; let gate_options = self .options .scoped_to_q8_k_activation_role(Q8KActivationMatvecRole::FfnGate); @@ -498,7 +504,7 @@ impl<'a> ScalarLlamaSession<'a> { let (mut gate, up) = if profile_events.is_none() && comparison_events.is_none() { let paired = layer.ffn_gate.mul_vec_pair_with_options( &layer.ffn_up, - &ffn_normed, + &normed, gate_options, up_options, )?; @@ -506,19 +512,15 @@ impl<'a> ScalarLlamaSession<'a> { pair } else { let (gate, up) = rayon::join( - || { - layer - .ffn_gate - .mul_vec_with_options(&ffn_normed, gate_options) - }, - || layer.ffn_up.mul_vec_with_options(&ffn_normed, up_options), + || layer.ffn_gate.mul_vec_with_options(&normed, gate_options), + || layer.ffn_up.mul_vec_with_options(&normed, up_options), ); (gate?, up?) } } else { let gate = profiled_layer_mul_vec( &layer.ffn_gate, - &ffn_normed, + &normed, layer_index, "ffn_gate", profile_events.as_deref_mut(), @@ -527,7 +529,7 @@ impl<'a> ScalarLlamaSession<'a> { )?; let up = profiled_layer_mul_vec( &layer.ffn_up, - &ffn_normed, + &normed, layer_index, "ffn_up", profile_events.as_deref_mut(), @@ -552,10 +554,11 @@ impl<'a> ScalarLlamaSession<'a> { let (token_id, logits) = match output_mode { OutputMode::Logits => { - let normed = rms_norm( + rms_norm_into( &hidden, &self.model.weights.output_norm, self.model.config.rms_norm_epsilon, + &mut normed, )?; let output = self .model @@ -574,10 +577,11 @@ impl<'a> ScalarLlamaSession<'a> { (Some(argmax(&logits)?), Some(logits)) } OutputMode::TokenIdOnly => { - let normed = rms_norm( + rms_norm_into( &hidden, &self.model.weights.output_norm, self.model.config.rms_norm_epsilon, + &mut normed, )?; let output = self .model diff --git a/docs/benchmarks/2026-07-17-reusable-normalization-scratch.md b/docs/benchmarks/2026-07-17-reusable-normalization-scratch.md new file mode 100644 index 0000000..283f8ea --- /dev/null +++ b/docs/benchmarks/2026-07-17-reusable-normalization-scratch.md @@ -0,0 +1,75 @@ +# Reusable normalization scratch diagnostic + +Date: 2026-07-17 + +## Scope + +This diagnostic evaluates one bounded single-stream allocation change. A token +step previously created a new root-mean-square normalization vector for the +attention input and feed-forward input of every transformer layer, then created +one more for the output projection. The candidate keeps one hidden-size vector +and refills it at each of those non-overlapping stages. + +For a model with `L` transformer layers, a token step now makes one +normalization-buffer allocation instead of `2L + 1`. Context-only prompt steps +avoid the final output normalization in both implementations. The arithmetic +order inside normalization is unchanged. + +## Fixed inputs + +- Host: Apple M5 Pro, 15 logical CPUs +- OS: macOS 26.5.2 arm64 +- Toolchain: Rust 1.96.0, LLVM 22.1.2 +- Source base: commit `8899f3ea61111597826ae773806073d4b9c8193f` + on `main`, with the candidate evaluated in a dirty working tree +- Cargo profile: repository `release` profile, no `RUSTFLAGS` +- Baseline CLI SHA-256: + `e286a8f5d7d3beb5e0660cb4928230b0a8905934a142bae39b69fef4ccc314c5` +- Candidate CLI SHA-256: + `ab158ca67378a5b73bfad2f4dc3083769b0b8c8525f7efebb5ea66e0eae9d88c` +- Model: Qwen2.5 0.5B Instruct Q4_K_M, 491,400,032 bytes +- Model SHA-256: + `74a4da8c9fdbcd15bd1f6d01d621410d31c6fc00986f5eb687824e7b93d7a9db` +- Prompt: `Write a short story about a rusty robot who learns to sail.` +- Measured decode steps: 128 per run +- Repetitions: five interleaved baseline and candidate pairs +- Workers: seven +- Policy: experimental residual Q8 activation matvec on Arm I8MM + +The host had unrelated interactive load. RSS and CPU were not sampled. The +timing result is therefore diagnostic, not clean-host release evidence. + +## Command + +Each copied release binary was run with the same working directory and model: + +```sh + \ + --model target/models/qwen2.5-0.5b-instruct-q4_k_m.gguf \ + --prompt 'Write a short story about a rusty robot who learns to sail.' \ + --benchmark-runs 128 \ + --threads 7 \ + --experimental-residual-q8-activation-matvec +``` + +The baseline and candidate were alternated for each pair. The complete +`benchmark_token_ids` line from every run was hashed separately from timing. + +## Result + +| Variant | Median decode step | Median decode rate | Token trace SHA-256 | +| --- | ---: | ---: | --- | +| Baseline | 10.373 ms | 96.41 tok/s | `ff342fb6b38a5f301d61dab6af424615a33dce9cc75ff1e9c026a2b40aa3674a` | +| Reusable scratch | 10.314 ms | 96.95 tok/s | `ff342fb6b38a5f301d61dab6af424615a33dce9cc75ff1e9c026a2b40aa3674a` | + +The candidate median was 0.57% faster. Pair-level timing was noisy, while all +ten runs produced the same complete ordered 128-token trace. The focused unit +test also verifies that repeated normalization writes retain the same vector +allocation and capacity. + +## Acceptance + +Accepted as an allocation reduction with exact trace parity and no API change. +The modest timing direction supports the change but is not promoted as a +clean-host throughput milestone. CPU, RSS, time to first token, and tail +latency claims require a future full eval on a quiet host. diff --git a/docs/benchmarks/README.md b/docs/benchmarks/README.md index 56042f9..370737b 100644 --- a/docs/benchmarks/README.md +++ b/docs/benchmarks/README.md @@ -11,6 +11,9 @@ experiments that still inform Ferrite's current architecture. - [`2026-07-14 bounded embedding row decode`](2026-07-14-bounded-embedding-row-decode.md) records the rejected whole-matrix decode, the bounded block-window fix, exact trace retention, and the diagnostic-only memory result. +- [`2026-07-17 reusable normalization scratch`](2026-07-17-reusable-normalization-scratch.md) + records the single-stream allocation reduction, exact 128-token trace + parity, and a contaminated-host interleaved timing diagnostic. - [`2026-07-09-concurrent-serving-phase1.md`](2026-07-09-concurrent-serving-phase1.md) records concurrent serving and batching evidence. - [`2026-07-10-oss-quality-hardening.md`](2026-07-10-oss-quality-hardening.md)