From 9f9b1e6f5a17db8f309ab61ce3e9abe8bf734ef2 Mon Sep 17 00:00:00 2001 From: EffortlessSteven <15812269+EffortlessSteven@users.noreply.github.com> Date: Sun, 21 Jun 2026 01:23:54 +0000 Subject: [PATCH] =?UTF-8?q?=E2=9A=A1=20Bolt:=20Optimize=20repetition=20pen?= =?UTF-8?q?alty=20hot=20path?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .jules/bolt.md | 4 ++++ crates/bitnet-sampling/src/strategies.rs | 10 ++++++++-- ...t_wave12__repetition_penalty_apply_with_counts.snap | 2 +- 3 files changed, 13 insertions(+), 3 deletions(-) diff --git a/.jules/bolt.md b/.jules/bolt.md index 2466a6e27d..5a5f7dbca1 100644 --- a/.jules/bolt.md +++ b/.jules/bolt.md @@ -5,3 +5,7 @@ ## 2026-03-05 - Hot Loop Allocations in Token Sampling **Learning:** Allocating memory in the hot path of LLM token generation (e.g., `logits.to_vec()` or creating `HashMap`s per token) significantly degrades performance due to repeated allocation overhead of vocabulary-sized vectors (often 128K+ elements). Additionally, mathematically equivalent iterative multiplication (`logit *= inv_penalty`) can replace `HashMap` counting and `.powi(count)`, completely eliminating O(N) memory allocations per token. **Action:** When working on generation loops, use buffer pooling (e.g. storing a `Vec` in the generator state and using `std::mem::take` to bypass borrow checker limitations) and avoid `HashMap` allocations for simple counting if an iterative scalar approach is mathematically equivalent. + +## 2026-03-05 - Repetition Penalty Power Calculation Optimization +**Learning:** Calculating `powi` unconditionally before an `if/else` block introduces significant overhead by performing work for branches that aren't taken. Furthermore, executing division (e.g., `logit /= penalty`) in a hot loop is computationally expensive and can be safely eliminated by calculating the inverse first (`1.0 / penalty`) and multiplying with the inverse `inv_penalty.powi(count)`. +**Action:** When working on math operations in conditional branches inside hot loops, pre-calculate values if possible outside the loop and lazily calculate values inside the branches that actually use them to avoid wasted cycles and unnecessary expensive operations like division. diff --git a/crates/bitnet-sampling/src/strategies.rs b/crates/bitnet-sampling/src/strategies.rs index 8cbc5fa903..5ada35e2d2 100644 --- a/crates/bitnet-sampling/src/strategies.rs +++ b/crates/bitnet-sampling/src/strategies.rs @@ -264,6 +264,9 @@ impl RepetitionPenaltyConfig { /// /// `token_counts` is a slice of `(token_id, occurrence_count)` pairs. pub fn apply(&self, logits: &mut [f32], token_counts: &[(u32, usize)]) { + // ⚡ Bolt: Hoist the inverse calculation entirely outside the loop to avoid redundant operations + let inv_count_penalty = 1.0 / self.count_penalty; + for &(token_id, count) in token_counts { let idx = token_id as usize; if idx >= logits.len() || count == 0 { @@ -279,11 +282,14 @@ impl RepetitionPenaltyConfig { // Count penalty: multiplicative if self.count_penalty.to_bits() != 1.0f32.to_bits() { let count = i32::try_from(count).unwrap_or(i32::MAX); - let penalty = self.count_penalty.powi(count); if logits[idx] > 0.0 { - logits[idx] /= penalty; + // ⚡ Bolt: Lazily compute inverse power inside branch and multiply to safely eliminate division overhead + let inv_penalty = inv_count_penalty.powi(count); + logits[idx] *= inv_penalty; } else { + // ⚡ Bolt: Lazily compute power inside branch to avoid unnecessary work when this branch is not taken + let penalty = self.count_penalty.powi(count); logits[idx] *= penalty; } } diff --git a/crates/bitnet-sampling/tests/snapshots/snapshot_wave12__repetition_penalty_apply_with_counts.snap b/crates/bitnet-sampling/tests/snapshots/snapshot_wave12__repetition_penalty_apply_with_counts.snap index 5cd68c115e..1b28c27544 100644 --- a/crates/bitnet-sampling/tests/snapshots/snapshot_wave12__repetition_penalty_apply_with_counts.snap +++ b/crates/bitnet-sampling/tests/snapshots/snapshot_wave12__repetition_penalty_apply_with_counts.snap @@ -6,6 +6,6 @@ expression: logits 1.0, 0.15026294, 3.0, - 2.909091, + 2.9090908, 5.0, ]