From 3b167ac8edf2cb99dbb1a66332aeb0815f91621f Mon Sep 17 00:00:00 2001 From: "The gemma.cpp Authors" Date: Sat, 25 Jul 2026 22:40:37 -0700 Subject: [PATCH] Optimize Gemma 4 MoE router input buffering and scaling: - Changed `router_in` storage type from `BF16` to `float` in `Activations` to avoid `BF16`->`float`->`BF16` round-tripping. - Copies pre-FFW activations directly to the `float` `router_in` without type conversion. - Performs fused inline scaling using Highway SIMD (`Decompress1AndCompressInplace`) avoiding per-token type roundtrips and temporary vector allocations. PiperOrigin-RevId: 954067115 --- gemma/activations.h | 2 +- gemma/gemma4_moe.cc | 42 ++++++++++++++++++------------------------ 2 files changed, 19 insertions(+), 25 deletions(-) diff --git a/gemma/activations.h b/gemma/activations.h index 085104b5..aba4541e 100644 --- a/gemma/activations.h +++ b/gemma/activations.h @@ -715,7 +715,7 @@ struct Activations { TensorStats s_ffw_out; // For MoE layers. These are used outside the expert-parallel loop: - MatStorageT router_in; + MatStorageT router_in; MatStorageT router_logits; // batch_size x num_experts // DeepSeek MLA (zero-sized unless a layer uses MLA). diff --git a/gemma/gemma4_moe.cc b/gemma/gemma4_moe.cc index a67d4e04..8e25e9b2 100644 --- a/gemma/gemma4_moe.cc +++ b/gemma/gemma4_moe.cc @@ -111,34 +111,28 @@ struct Gemma4MoE { for (size_t token_idx = 0; token_idx < num_tokens; ++token_idx) { const float* pre_ffw_row = activations.x.Row(token_idx); - BF16* router_in_row = activations.router_in.Row(token_idx); - - for (size_t col = 0; col < model_dim; ++col) { - router_in_row[col] = hwy::ConvertScalarTo(pre_ffw_row[col]); - } + float* router_in_row = activations.router_in.Row(token_idx); + std::copy_n(pre_ffw_row, model_dim, router_in_row); } RMSNormNoScaleInplaceBatched(activations.router_in, env.ctx); - // TODO(philculliton): Use a float buffer for router_in to avoid the - // BF16->float->BF16 round-trip, and precompute scale_factor * router_scale - // once rather than per-token. Per the CL comment: we are converting to - // bf16, but then converting back to float below. Should we set up a - // router_in_row_f32 so we can just keep it as float? (That would help if - // num_tokens>>1, because we could precompute * scale_factor once.) - for (size_t token_idx = 0; token_idx < num_tokens; ++token_idx) { - BF16* router_in_row = activations.router_in.Row(token_idx); - if (has_router_scale) { - for (size_t col = 0; col < model_dim; ++col) { - router_in_row[col] = hwy::ConvertScalarTo( - hwy::ConvertScalarTo(router_in_row[col]) * scale_factor * - hwy::ConvertScalarTo(scale_ptr[col])); - } - } else { - for (size_t col = 0; col < model_dim; ++col) { - router_in_row[col] = hwy::ConvertScalarTo( - hwy::ConvertScalarTo(router_in_row[col]) * scale_factor); - } + namespace hn = hwy::HWY_NAMESPACE; + using DF = hn::ScalableTag; + + if (has_router_scale) { + for (size_t token_idx = 0; token_idx < num_tokens; ++token_idx) { + float* router_in_row = activations.router_in.Row(token_idx); + Decompress1AndCompressInplace( + DF(), router_in_row, model_dim, scale_ptr, 0, + [scale_factor](auto df, auto inout, auto scale) HWY_ATTR { + return hn::Mul(inout, hn::Mul(scale, hn::Set(df, scale_factor))); + }); + } + } else { + for (size_t token_idx = 0; token_idx < num_tokens; ++token_idx) { + float* router_in_row = activations.router_in.Row(token_idx); + MulByConst(scale_factor, router_in_row, model_dim); } }