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); } }