Reuse normalization scratch buffer

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

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<Vec<f32>, 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<f32>,
+) -> Result<(), InferenceError> {
     if input.is_empty() {
         return Err(InferenceError::new("rms_norm input must not be empty"));
     }
@@ -32,15 +43,16 @@ pub fn rms_norm(input: &[f32], weight: &[f32], epsilon: f32) -> Result<Vec<f32>,
         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() {
             return Err(InferenceError::new("rms_norm output must be finite"));
         }
         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];

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,27 +504,23 @@ 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,
                 )?;
                 if let Some(pair) = paired {
                     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

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
+<baseline-or-candidate-binary> \
+  --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.

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)