From 3b270d236fc1ed73493ed24134e6c137f460cab9 Mon Sep 17 00:00:00 2001 From: Nanubala Gnana Sai <45007169+jonpsy@users.noreply.github.com> Date: Tue, 29 Oct 2024 22:38:05 +0530 Subject: [PATCH] Use hwy divisor to speed up division --- gemma/gemma-inl.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/gemma/gemma-inl.h b/gemma/gemma-inl.h index ea5aca0..9fd739a 100644 --- a/gemma/gemma-inl.h +++ b/gemma/gemma-inl.h @@ -319,7 +319,7 @@ class GemmaAttention { head * layer_config_.qkv_dim * 2; KVCache& kv_cache = kv_caches_[query_idx]; - const size_t grp_size = layer_config_.grp_size; + const hwy::Divisor& div_grp_size { static_cast(layer_config_.grp_size) }; const size_t ngb_size = layer_config_.ngb_size; const bool self_extend = layer_config_.self_extend; @@ -330,7 +330,7 @@ class GemmaAttention { // When embedding position, we will use grouped key position if (self_extend && pos > ngb_size) { - pos /= grp_size; + pos = div_grp_size.Divide(pos); } // Copy from `q` if MHA, or apply in-place. PositionalEncodingQK(is_mha_ ? mha_kv : kv, pos, layer_, 1.0f,