From 5df0438a398c53ebdff72fdcf667d18fd6a1c90d Mon Sep 17 00:00:00 2001 From: Phil Culliton Date: Mon, 13 Jul 2026 13:06:39 -0700 Subject: [PATCH] Gemma 4 E2B text model PiperOrigin-RevId: 947201464 --- gemma/activations.h | 26 +++++- gemma/attention.cc | 3 +- gemma/configs.cc | 65 ++++++++++++- gemma/configs.h | 14 ++- gemma/flash_attention.cc | 2 +- gemma/gemma-inl.h | 2 + gemma/gemma.cc | 193 ++++++++++++++++++++++++++++++++++++++- gemma/gemma4_moe.cc | 11 ++- gemma/kv_cache.cc | 82 ++++++++++++++--- gemma/tensor_info.cc | 57 ++++++++++++ gemma/tiled_attention.cc | 81 ++++++++-------- gemma/tokenizer.cc | 4 + gemma/vit.cc | 1 + gemma/weights.h | 27 ++++++ python/configs.cc | 1 + util/mat.h | 2 +- 16 files changed, 499 insertions(+), 72 deletions(-) diff --git a/gemma/activations.h b/gemma/activations.h index 481926f2..a4c2576d 100644 --- a/gemma/activations.h +++ b/gemma/activations.h @@ -47,6 +47,14 @@ static inline size_t MaxQkvDim(const ModelConfig& config) { } return max_dim; } +static inline size_t MaxFFHiddenDim(const ModelConfig& config) { + size_t max_dim = 0; + for (const auto& lc : config.layer_configs) { + max_dim = HWY_MAX(max_dim, static_cast(lc.ff_hidden_dim)); + } + return max_dim; +} + static inline float ChooseQueryScale(const ModelConfig& config) { const LayerConfig& layer_config = config.layer_configs[0]; @@ -115,9 +123,7 @@ struct AttentionActivations { layer_config.post_qk == PostQKType::HalfRope)), inv_timescale_global(CreateInvTimescale( allocator, - config.partial_rotary_factor < 1.0f - ? max_qkv_dim - : max_qkv_dim / 4, + max_qkv_dim, layer_config.post_qk == PostQKType::HalfRope, 1000000.0, config.partial_rotary_factor)) { // Batch size can be 0 in experimental code so do not assert. @@ -360,12 +366,15 @@ struct Activations { pre_ffw_rms_out(MatFactory("pre_ffw_rms_out", batch_size, config.model_dim, ctx.allocator)), - C1(MatFactory("C1", batch_size, layer_config.ff_hidden_dim, + C1(MatFactory("C1", batch_size, MaxFFHiddenDim(config), ctx.allocator)), - C2(MatFactory("C2", batch_size, layer_config.ff_hidden_dim, + C2(MatFactory("C2", batch_size, MaxFFHiddenDim(config), ctx.allocator)), ffw_out( MatFactory("ffw_out", batch_size, config.model_dim, ctx.allocator)), + ple_embeds( + MatFactory("ple_embeds", batch_size, + config.num_layers * config.ple_dim, ctx.allocator)), max_workers(ctx.pools.MaxWorkers()), s_ffw_in(config.num_layers, max_workers), @@ -401,6 +410,9 @@ struct Activations { x.AllocateAndAttachRowPtrs(row_ptrs); x_bf.AllocateAndAttachRowPtrs(row_ptrs); logits.AllocateAndAttachRowPtrs(row_ptrs); + if (config.ple_dim > 0) { + ple_embeds.AllocateAndAttachRowPtrs(row_ptrs); + } C1.AllocateAndAttachRowPtrs(row_ptrs); C2.AllocateAndAttachRowPtrs(row_ptrs); ffw_out.AllocateAndAttachRowPtrs(row_ptrs); @@ -461,6 +473,9 @@ struct Activations { C1.OverrideRows(batch_size); C2.OverrideRows(batch_size); ffw_out.OverrideRows(batch_size); + if (layer_config.ple_dim > 0) { + ple_embeds.OverrideRows(batch_size); + } attention_storage.SetBatchSize(batch_size); // `AttentionActivationsPtrs` holds `MatPtrT` which also require updating; @@ -480,6 +495,7 @@ struct Activations { MatStorageT C1; MatStorageT C2; MatStorageT ffw_out; + MatStorageT ple_embeds; const size_t max_workers; TensorStats s_ffw_in; diff --git a/gemma/attention.cc b/gemma/attention.cc index c608e8ac..9ec3d93f 100644 --- a/gemma/attention.cc +++ b/gemma/attention.cc @@ -198,6 +198,7 @@ static HWY_INLINE void ComputeQKV(size_t num_tokens, const size_t layer_idx, CallMatMul(activations.pre_att_rms_out, layer.qkv_einsum_w1, /*add=*/nullptr, env, activations.q); + if (flags & kSkipKV) return; // Set up MatMul row pointers for writing to KV, which consists of // `kv_heads` pairs of (k, v) vectors. This safely handles wraparound // because rows are computed modulo seq_len. @@ -294,7 +295,7 @@ static HWY_INLINE void ComputeQKV(size_t num_tokens, const size_t layer_idx, RMSNormInplace(weights_t->PackedScale1(), /*w_ofs=*/0, kv_f32, qkv_dim, env.ctx, worker); }); - } else if (layer_config.post_qk == PostQKType::NormLocalRope) { + } else if (layer_config.post_qk == PostQKType::NormLocalRope || layer_config.use_qk_norm) { RMSNormNoScaleInplace(kv_f32, qkv_dim, env.ctx, worker); } diff --git a/gemma/configs.cc b/gemma/configs.cc index 833c3142..09a1f598 100644 --- a/gemma/configs.cc +++ b/gemma/configs.cc @@ -470,7 +470,7 @@ static LayerConfig LayerConfigGemma4_26B_MoE_LM(size_t model_dim) { static ModelConfig ConfigGemma4_26B_MoE() { ModelConfig config = ConfigBaseGemmaV4(); config.display_name = "Gemma4_26B_MoE"; - config.final_cap = 30.0f; + config.final_cap = 0.0f; config.att_cap = 0.0f; config.model = Model::GEMMA4_26B_MOE; config.wrapping = PromptWrapping::GEMMA_IT; @@ -495,6 +495,62 @@ static ModelConfig ConfigGemma4_26B_MoE() { return config; } +static LayerConfig LayerConfigGemma4_2B_Local(size_t model_dim) { + LayerConfig config; + config.model_dim = model_dim; + config.ff_hidden_dim = 6144; + config.heads = 8; + config.kv_heads = 1; + config.qkv_dim = 256; + config.optimized_gating = true; + config.post_norm = PostNormType::Scale; + config.activation = ActivationType::Gelu; + config.post_qk = PostQKType::NormLocalRope; + config.use_qk_norm = true; + config.norm_v = true; + config.ple_dim = 256; + return config; +} + +static LayerConfig LayerConfigGemma4_2B_Global(size_t model_dim) { + LayerConfig config = LayerConfigGemma4_2B_Local(model_dim); + config.qkv_dim = 512; + return config; +} + +// Until we have the audio checkpoints included, we use the LM config directly. +static ModelConfig ConfigGemma4_2B() { + ModelConfig config = ConfigBaseGemmaV4(); + config.display_name = "Gemma4_2B"; + config.model = Model::GEMMA4_2B; + config.wrapping = PromptWrapping::GEMMA_IT; + config.model_dim = 1536; + config.vocab_size = kGemmaV3VocabSize; // 262144 + config.max_seq_len = 128 * 1024; + config.final_cap = 0.0f; + config.ple_dim = 256; + config.num_layers = 35; + config.use_global_timescale = true; + config.partial_rotary_factor = 0.25f; + config.query_scale = QueryScaleType::One; + LayerConfig local_config = LayerConfigGemma4_2B_Local(config.model_dim); + config.layer_configs = {config.num_layers, local_config}; + // Global attention layers: [4, 9, 14, 19, 24, 29, 34] (stride 5) + for (size_t i = 0; i < config.num_layers; ++i) { + if (i % 5 == 4) { + config.layer_configs[i] = LayerConfigGemma4_2B_Global(config.model_dim); + } + } + // Double-wide MLP for last 20 layers (KV-shared layers 15-34) + for (size_t i = 15; i < config.num_layers; ++i) { + config.layer_configs[i].ff_hidden_dim = 12288; + config.layer_configs[i].kv_share_layer_idx = (i % 5 == 4) ? 14 : 13; + } + config.attention_window_sizes = RepeatedAttentionWindowSizes<35, 5>( + {512, 512, 512, 512, config.max_seq_len}); + return config; +} + static ModelConfig ConfigFromModel(Model model) { switch (model) { case Model::GEMMA2_2B: @@ -529,6 +585,8 @@ static ModelConfig ConfigFromModel(Model model) { return ConfigGemma3_27B_LM(); case Model::GEMMA4_26B_MOE: return ConfigGemma4_26B_MoE(); + case Model::GEMMA4_2B: + return ConfigGemma4_2B(); default: HWY_ABORT("Model type %d unknown.", static_cast(model)); } @@ -570,6 +628,8 @@ const char* ModelPrefix(Model model) { return "gemma3-27b-lm"; case Model::GEMMA4_26B_MOE: return "gemma4-26b-moe"; + case Model::GEMMA4_2B: + return "gemma4-2b"; default: HWY_ABORT("Model type %d unknown.", static_cast(model)); } @@ -784,7 +844,8 @@ Model DeduceModel(const Path& blob_path, size_t layers, int layer_types) { case 62: return (layer_types & kDeducedViT) ? Model::GEMMA3_27B : Model::GEMMA3_27B_LM; - + case 35: + return Model::GEMMA4_2B; // TODO: detect these. /* return Model::GEMMA2_772M; diff --git a/gemma/configs.h b/gemma/configs.h index a4d270d7..38809b6e 100644 --- a/gemma/configs.h +++ b/gemma/configs.h @@ -35,6 +35,8 @@ namespace gcpp { constexpr size_t kMaxBF16PerVector = HWY_ARCH_MAX_BYTES / sizeof(BF16); +HWY_INLINE_VAR constexpr int kSkipKV = 1; + HWY_INLINE_VAR constexpr size_t kMaxQKVDim = 1024; #ifndef GEMMA_FUSED_FFN @@ -131,7 +133,7 @@ enum class PostQKType { Rope, HalfRope, NormLocalRope = 8, // Norm without scale, and rope for local attention layers - kSentinel // must be last + kSentinel // must be last }; static inline bool EnumValid(PostQKType type) { @@ -221,6 +223,7 @@ enum class Model { GEMMA3_12B_LM, GEMMA3_27B_LM, GEMMA4_26B_MOE, + GEMMA4_2B, kSentinel, }; @@ -303,6 +306,8 @@ struct LayerConfig : public IFields { visitor(norm_v); visitor(num_experts); visitor(num_experts_per_datapoint); + visitor(ple_dim); + visitor(kv_share_layer_idx); // Append new fields here, then update `python/configs.cc`. } @@ -338,6 +343,8 @@ struct LayerConfig : public IFields { bool norm_v = false; // Normalize V projections before caching. uint32_t num_experts = 0; uint32_t num_experts_per_datapoint = 0; + uint32_t ple_dim = 0; // Per-Layer Embedding dimension (0 = disabled). + int kv_share_layer_idx = -1; InternalLayerConfig internal; }; @@ -437,6 +444,7 @@ struct ModelConfig : public IFields { visitor(use_global_timescale); visitor(partial_rotary_factor); + visitor(ple_dim); // Append new fields here, then update `python/configs.cc`. } @@ -552,7 +560,9 @@ struct ModelConfig : public IFields { InternalModelConfig internal; bool use_global_timescale = false; // for Gemma 3 - float partial_rotary_factor = 1.0f; // Fraction of dims with RoPE (0.25 for Gemma4 MoE). + float partial_rotary_factor = + 1.0f; // Fraction of dims with RoPE (0.25 for Gemma4 MoE). + uint32_t ple_dim = 0; // Per-Layer Embedding dimension (0 = disabled). }; // Returns the sub-config for the ViT model of the PaliGemma model. diff --git a/gemma/flash_attention.cc b/gemma/flash_attention.cc index c591c563..1383f0d1 100644 --- a/gemma/flash_attention.cc +++ b/gemma/flash_attention.cc @@ -91,7 +91,7 @@ void RMSNormAndPositionalEncoding(const size_t num_tokens, const QBatch& qbatch, RMSNormInplace(weights_t->PackedScale1(), /*w_ofs=*/0, q_row, layer_config.qkv_dim, ctx, worker); }); - } else if (layer_config.post_qk == PostQKType::NormLocalRope) { + } else if (layer_config.post_qk == PostQKType::NormLocalRope || layer_config.use_qk_norm) { RMSNormNoScaleInplace(q_row, layer_config.qkv_dim, ctx, worker); } PositionalEncodingQK(q_row, layer_idx, activations, ctx, worker, pos, diff --git a/gemma/gemma-inl.h b/gemma/gemma-inl.h index 91f22752..2035c961 100644 --- a/gemma/gemma-inl.h +++ b/gemma/gemma-inl.h @@ -177,6 +177,8 @@ static inline void FFWNoVit(const LayerWeightsPtrs& layer, Activations& activations, MatMulEnv& env) { GCPP_ZONE(env.ctx, hwy::Profiler::GlobalIdx(), Zones::kGenFFW); const LayerConfig& layer_config = layer.layer_config; + activations.C1.OverrideCols(layer_config.ff_hidden_dim); + activations.C2.OverrideCols(layer_config.ff_hidden_dim); HWY_DASSERT(!layer_config.ff_biases); // Only used in Vit. diff --git a/gemma/gemma.cc b/gemma/gemma.cc index d55dc3dd..899dd4f6 100644 --- a/gemma/gemma.cc +++ b/gemma/gemma.cc @@ -105,16 +105,28 @@ static HWY_NOINLINE void TransformerLayer(const size_t num_tokens, const LayerConfig& layer_config = layer.layer_config; if (layer_config.IsMoE() && activations.attention.config.model == Model::GEMMA4_26B_MOE) { - Gemma4MoETransformerLayer(num_tokens, layer_idx, layer, activations, - qbatch, env); + Gemma4MoETransformerLayer(num_tokens, layer_idx, layer, activations, qbatch, + env); return; } RMSNormBatched(activations.x, layer.pre_attention_norm_scale, activations.attention.pre_att_rms_out, env.ctx); - Attention(layer_config.type, num_tokens, layer_idx, layer, activations, - qbatch, env); + // gemma.cc -> TransformerLayer + const size_t kv_cache_layer_idx = + (layer_config.kv_share_layer_idx >= 0) + ? static_cast(layer_config.kv_share_layer_idx) + : layer_idx; + + const int flags = (kv_cache_layer_idx == layer_idx) ? 0 : kSkipKV; + if (layer_config.type == LayerAttentionType::kGemma) { + GemmaAttention(num_tokens, kv_cache_layer_idx, layer, activations.attention, + qbatch, env, activations.attention_impl, flags); + } else { + Attention(layer_config.type, num_tokens, layer_idx, layer, activations, + qbatch, env); + } PostNorm(layer_config.post_norm, layer.post_attention_norm_scale, activations.attention.att_sums, env.ctx); @@ -136,6 +148,87 @@ static HWY_NOINLINE void TransformerLayer(const size_t num_tokens, ResidualConnection(activations.ffw_out, activations.x, layer, /*is_attention=*/false, env.ctx); + if (layer_config.ple_dim > 0) { + // 1. Gate: [batch, model_dim] @ [model_dim, ple_dim] -> [batch, ple_dim] + // Use activations.x_bf to convert activations.x + for (size_t r = 0; r < num_tokens; ++r) { + for (size_t c = 0; c < layer_config.model_dim; ++c) { + activations.x_bf.Row(r)[c] = BF16(activations.x.Row(r)[c]); + } + } + + MatStorageT gate_out(MatFactory( + "gate_out", num_tokens, layer_config.ple_dim, env.ctx.allocator)); + std::vector> row_ptrs; + gate_out.AllocateAndAttachRowPtrs(row_ptrs); + + CallMatMul(activations.x_bf, layer.ple_gate, /*add=*/nullptr, env, + gate_out); + + // 2. Activation and Element-wise multiply + const size_t ple_dim = layer_config.ple_dim; + const size_t layer_offset = layer_idx * ple_dim; + + ParallelFor(Parallelism::kFlat, num_tokens, env.ctx, /*cluster_idx=*/0, + Callers::kActivationBatched, + [&](uint64_t token_idx, size_t worker) { + float* g_row = gate_out.Row(token_idx); + const float* p_row = + activations.ple_embeds.Row(token_idx) + layer_offset; + namespace hn = hwy::HWY_NAMESPACE; + using DF = hn::ScalableTag; + using VF = hn::Vec; + const DF df; + for (size_t i = 0; i < ple_dim; i += hn::Lanes(df)) { + VF v = hn::Load(df, g_row + i); + v = Gelu(df, v); + VF p = hn::Load(df, p_row + i); + v = hn::Mul(v, p); + hn::Store(v, df, g_row + i); + } + }); + + // 3. Projection: [batch, ple_dim] @ [ple_dim, model_dim] -> [batch, + // model_dim] + MatStorageT gate_out_bf(MatFactory( + "gate_out_bf", num_tokens, layer_config.ple_dim, env.ctx.allocator)); + gate_out_bf.AllocateAndAttachRowPtrs(row_ptrs); + for (size_t r = 0; r < num_tokens; ++r) { + for (size_t c = 0; c < ple_dim; ++c) { + gate_out_bf.Row(r)[c] = BF16(gate_out.Row(r)[c]); + } + } + + CallMatMul(gate_out_bf, layer.ple_proj, /*add=*/nullptr, env, + activations.ffw_out); + + // 4. Norm and Residual Add + RMSNormInplaceBatched(layer.post_ple_ns, activations.ffw_out, env.ctx); + + ParallelFor(Parallelism::kFlat, num_tokens, env.ctx, /*cluster_idx=*/0, + Callers::kOpsAddFromBatched, + [&](uint64_t token_idx, size_t worker) { + AddFrom(activations.ffw_out.Row(token_idx), + activations.x.Row(token_idx), layer_config.model_dim, + env.ctx, worker); + }); + } + + if (layer.skip_scale.HasPtr()) { + float skip_scale_val = 1.0f; + if (layer.skip_scale.GetType() == Type::kF32) { + skip_scale_val = static_cast(layer.skip_scale.Packed())[0]; + } else if (layer.skip_scale.GetType() == Type::kBF16) { + skip_scale_val = hwy::ConvertScalarTo( + static_cast(layer.skip_scale.Packed())[0]); + } else { + HWY_ABORT("Unexpected skip_scale type: %d", + static_cast(layer.skip_scale.GetType())); + } + for (size_t r = 0; r < activations.x.Rows(); ++r) { + MulByConst(skip_scale_val, activations.x.Row(r), activations.x.Cols()); + } + } } // Returns the scale value to use for the embedding (basically sqrt model_dim). @@ -205,6 +298,73 @@ EmbedMMToken(int token, size_t x_row, size_t pos, size_t pos_in_prompt, return image_token_position; } +static HWY_NOINLINE void ComputePLEEmbeddings(size_t tbatch_size, + const std::vector& tokens, + const ModelConfig& config, + const WeightsPtrs& weights, + Activations& activations, + MatMulEnv& env) { + if (config.ple_dim == 0) return; + + // 1. Convert activations.x (float) to activations.x_bf (BF16) + for (size_t r = 0; r < tbatch_size; ++r) { + for (size_t c = 0; c < config.model_dim; ++c) { + activations.x_bf.Row(r)[c] = BF16(activations.x.Row(r)[c]); + } + } + + // 2. CallMatMul for the context projection + CallMatMul(activations.x_bf, weights.ple_model_proj, /*add=*/nullptr, env, + activations.ple_embeds); + + // 3. Apply the model projection scale + const float scale_proj = 1.0f / sqrtf(static_cast(config.model_dim)); + const size_t ple_total_dim = config.num_layers * config.ple_dim; + + for (size_t r = 0; r < tbatch_size; ++r) { + float* out_row = activations.ple_embeds.Row(r); + for (size_t c = 0; c < ple_total_dim; ++c) { + out_row[c] *= scale_proj; + } + } + + // 4. RMSNorm (applied to each layer's embedding independently) + CallUpcasted(&weights.ple_proj_norm, [&](const auto* weights_t) { + ParallelFor(Parallelism::kFlat, tbatch_size, env.ctx, /*cluster_idx=*/0, + Callers::kOpsRMSNormInplaceBatched, + [&](uint64_t token_idx, size_t worker) { + float* row = activations.ple_embeds.Row(token_idx); + for (size_t layer = 0; layer < config.num_layers; ++layer) { + float* slice = row + layer * config.ple_dim; + RMSNormInplace(weights_t->PackedScale1(), /*w_ofs=*/0, + slice, config.ple_dim, env.ctx, worker); + } + }); + }); + + // 5. Add token embedding and apply input scale + const float scale_input = 1.0f / sqrtf(2.0f); + std::vector token_emb(ple_total_dim); + for (size_t r = 0; r < tbatch_size; ++r) { + int token = tokens[r]; + CallUpcasted(&weights.ple_embeddings, [&](const auto* weights_t) { + const size_t embedding_ofs = token * weights_t->Stride(); + const auto embedding_span = + MakeSpan(weights_t->Row(0), embedding_ofs + ple_total_dim); + const hn::ScalableTag df; + DecompressAndZeroPad(df, embedding_span, embedding_ofs, token_emb.data(), + ple_total_dim); + + const float token_scale = + sqrtf(static_cast(config.ple_dim)) * weights_t->Scale(); + float* out_row = activations.ple_embeds.Row(r); + for (size_t c = 0; c < ple_total_dim; ++c) { + out_row[c] = (out_row[c] + token_emb[c] * token_scale) * scale_input; + } + }); + } +} + // Populates KV cache for batches of tokens from one query at a time. This is // called if prompts are longer than the query batch size, and also in // prefix-LM mode (end > 0), which must see all tokens in one batch. @@ -265,11 +425,18 @@ static HWY_NOINLINE void PrefillTBatch(const ModelConfig& config, // Fill activations.x (much faster than TransformerLayer). size_t image_token_position = 0; + std::vector tbatch_tokens; + if (config.ple_dim > 0) { + tbatch_tokens.reserve(tbatch_size); + } for (size_t ti = 0; ti < tbatch_size; ++ti) { const size_t pos = qbatch_1.Pos(0) + ti; const size_t pos_in_prompt = tbatch_start + ti; HWY_DASSERT(pos_in_prompt < prompt_size); const int token = qbatch_1.Prompt(0)[pos_in_prompt]; + if (config.ple_dim > 0) { + tbatch_tokens.push_back(token); + } image_token_position = EmbedMMToken( token, ti, pos, pos_in_prompt, config, weights, activations.x, env.ctx, runtime_config.image_tokens, image_token_position); @@ -282,6 +449,10 @@ static HWY_NOINLINE void PrefillTBatch(const ModelConfig& config, HWY_ASSERT(attend_to_last_token); } } + if (config.ple_dim > 0) { + ComputePLEEmbeddings(tbatch_size, tbatch_tokens, config, weights, + activations, env); + } // Transformer with one batch of tokens from a single query. No need to // set `PrevToken` because we already did the embedding above. @@ -336,10 +507,22 @@ static HWY_NOINLINE void Transformer(const ModelConfig& config, } // TODO: parallelize? + std::vector tbatch_tokens; + if (config.ple_dim > 0) { + tbatch_tokens.reserve(qbatch.Size()); + } for (size_t qi = 0; qi < qbatch.Size(); ++qi) { - EmbedMMToken(qbatch.PrevToken(qi), qi, qbatch.Pos(qi), + const int token = qbatch.PrevToken(qi); + if (config.ple_dim > 0) { + tbatch_tokens.push_back(token); + } + EmbedMMToken(token, qi, qbatch.Pos(qi), /*pos_in_prompt=*/0, config, weights, activations.x, env.ctx); } + if (config.ple_dim > 0) { + ComputePLEEmbeddings(qbatch.Size(), tbatch_tokens, config, weights, + activations, env); + } for (size_t layer_idx = 0; layer_idx < weights.c_layers.size(); ++layer_idx) { TransformerLayer(/*num_tokens=*/1, layer_idx, *weights.GetLayer(layer_idx), diff --git a/gemma/gemma4_moe.cc b/gemma/gemma4_moe.cc index e54efe95..311b0cca 100644 --- a/gemma/gemma4_moe.cc +++ b/gemma/gemma4_moe.cc @@ -507,8 +507,15 @@ void Gemma4MoETransformerLayer(size_t num_tokens, const size_t layer_idx, // Apply skip_scale AFTER residual connection // (HF: hidden_states *= self.layer_scalar, applied to full output) if (layer.skip_scale.HasPtr()) { - const float skip_scale_val = hwy::ConvertScalarTo( - *static_cast(layer.skip_scale.Packed())); + float skip_scale_val = 1.0f; + if (layer.skip_scale.GetType() == Type::kF32) { + skip_scale_val = static_cast(layer.skip_scale.Packed())[0]; + } else if (layer.skip_scale.GetType() == Type::kBF16) { + skip_scale_val = hwy::ConvertScalarTo( + static_cast(layer.skip_scale.Packed())[0]); + } else { + HWY_ABORT("Unexpected skip_scale type: %d", static_cast(layer.skip_scale.GetType())); + } for (size_t r = 0; r < num_tokens; ++r) { MulByConst(skip_scale_val, activations.x.Row(r), activations.x.Cols()); } diff --git a/gemma/kv_cache.cc b/gemma/kv_cache.cc index b7204aee..90718ca7 100644 --- a/gemma/kv_cache.cc +++ b/gemma/kv_cache.cc @@ -80,12 +80,60 @@ KVCache::KVCache(const Extents2D& kv_extents, size_t num_layers, // fires. The 2-arg constructor path relies on k_v_cols == 0 to skip reshape. } +// Support heterogeneous layer configurations (common in Gemma 4 architectures), +// where different layers can have varying attention shapes (e.g., mixing local +// layers with smaller qkv_dim/more heads and global layers with larger +// qkv_dim/fewer heads). +// +// Rather than assuming uniform layer sizes, we dynamically compute and store +// cumulative offsets for each layer to allow correct indexing into the +// flattened KV cache. KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, const Allocator& allocator) - : KVCache( - Extents2D(CappedSeqLen(config, inference_args), config.KVCacheCols()), - config.layer_configs.size(), config.layer_configs[0].kv_heads, - config.layer_configs[0].qkv_dim, allocator) {} + : allocator_(allocator) { + HWY_ASSERT(config.num_layers > 0); + HWY_ASSERT(config.num_layers == config.layer_configs.size()); + num_layers = config.num_layers; + + // 1. Build non-uniform offset tables dynamically + layer_flat_offsets.resize(num_layers, 0); + layer_k_v_offsets.resize(num_layers, 0); + rounded_qkv_dims.resize(num_layers, 0); + + size_t flat_accum = 0; + size_t k_v_accum = 0; + + for (size_t i = 0; i < num_layers; ++i) { + layer_flat_offsets[i] = static_cast(flat_accum); + flat_accum += config.layer_configs[i].CacheLayerSize(); + + layer_k_v_offsets[i] = static_cast(k_v_accum); + size_t rounded_dim = + hwy::RoundUpTo(config.layer_configs[i].qkv_dim, kMaxBF16PerVector); + rounded_qkv_dims[i] = static_cast(rounded_dim); + k_v_accum += config.layer_configs[i].kv_heads * rounded_dim; + } + k_v_cols = static_cast(k_v_accum); + + // Since we also store legacy homogeneous variables, we default them to Layer + // 0 values. + kv_heads = config.layer_configs[0].kv_heads; + qkv_dim = config.layer_configs[0].qkv_dim; + rounded_qkv_dim = hwy::RoundUpTo(qkv_dim, kMaxBF16PerVector); + + const size_t rows = CappedSeqLen(config, inference_args); + const size_t cols = config.KVCacheCols(); + kv_cache = MatStorageT("kv", Extents2D(rows, cols), allocator, + MatPadding::kOdd); + k_cache = MatStorageT( + "k", Extents2D(hwy::RoundUpTo(rows, kMaxBF16PerVector), k_v_cols), + allocator, MatPadding::kPacked); + v_cache = MatStorageT( + "v", Extents2D(hwy::RoundUpTo(rows, kMaxBF16PerVector), k_v_cols), + allocator, MatPadding::kPacked); + const size_t num_tiles = hwy::DivCeil(rows, kTileSize); + tiled_seq_len = num_tiles * kTileSize; +} KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, const RuntimeConfig& runtime_config, @@ -109,7 +157,8 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, flat_accum += config.layer_configs[i].CacheLayerSize(); layer_k_v_offsets[i] = static_cast(k_v_accum); - size_t rounded_dim = hwy::RoundUpTo(config.layer_configs[i].qkv_dim, kMaxBF16PerVector); + size_t rounded_dim = + hwy::RoundUpTo(config.layer_configs[i].qkv_dim, kMaxBF16PerVector); rounded_qkv_dims[i] = static_cast(rounded_dim); k_v_accum += config.layer_configs[i].kv_heads * rounded_dim; @@ -139,12 +188,14 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, allocator, MatPadding::kOdd); k_cache = MatStorageT( "k", - Extents2D(hwy::RoundUpTo(CappedSeqLen(config, inference_args), kMaxBF16PerVector), + Extents2D(hwy::RoundUpTo(CappedSeqLen(config, inference_args), + kMaxBF16PerVector), k_v_cols), allocator, MatPadding::kPacked); v_cache = MatStorageT( "v", - Extents2D(hwy::RoundUpTo(CappedSeqLen(config, inference_args), kMaxBF16PerVector), + Extents2D(hwy::RoundUpTo(CappedSeqLen(config, inference_args), + kMaxBF16PerVector), k_v_cols), allocator, MatPadding::kPacked); const size_t num_tiles = @@ -189,10 +240,10 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, size_t total_num_tiles = 0; for (size_t i = 0; i < num_layers; ++i) { - total_num_tiles += - num_tiles_per_head(config.attention_window_sizes[i], runtime_config.prefill_tbatch_size, - config.max_seq_len) * - config.layer_configs[i].kv_heads; + total_num_tiles += num_tiles_per_head(config.attention_window_sizes[i], + runtime_config.prefill_tbatch_size, + config.max_seq_len) * + config.layer_configs[i].kv_heads; } Extents2D extents(total_num_tiles, max_tile_length); compact_kv_cache_ptr = MatPtr("kv_tiled", kv_cache_type, extents); @@ -209,14 +260,15 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, kv_head_ptrs.clear(); kv_head_ptrs.reserve(num_layers * max_kv_heads); for (size_t i = 0; i < num_layers; ++i) { - size_t layer_tile_length = 2 * config.layer_configs[i].qkv_dim * kTileSize; + size_t layer_tile_length = + 2 * config.layer_configs[i].qkv_dim * kTileSize; if (kv_cache_type == Type::kInt8) { layer_tile_length += 2 * sizeof(BF16) * kTileSize; } for (size_t kv = 0; kv < config.layer_configs[i].kv_heads; ++kv) { - size_t num_tiles_per_kv_head = - num_tiles_per_head(config.attention_window_sizes[i], runtime_config.prefill_tbatch_size, - config.max_seq_len); + size_t num_tiles_per_kv_head = num_tiles_per_head( + config.attention_window_sizes[i], + runtime_config.prefill_tbatch_size, config.max_seq_len); MatPtr kv_ptr("kv_ptr", kv_cache_type, Extents2D(num_tiles_per_kv_head, layer_tile_length)); kv_ptr.SetPtr(compact_kv_cache_ptr.RowBytes(total_num_tiles), diff --git a/gemma/tensor_info.cc b/gemma/tensor_info.cc index 0d48b495..6635f98d 100644 --- a/gemma/tensor_info.cc +++ b/gemma/tensor_info.cc @@ -121,6 +121,39 @@ void TensorInfoRegistry::AddModelTensors(const ModelConfig& config) { .shape = {config.vit_config.model_dim}, .min_size = Type::kBF16, }); + // Per-Layer Embedding (PLE) model-level tensors. + if (config.ple_dim > 0) { + Add(no_suffix, + { + .base_name = "ple_embeddings", + .source_names = {"embedder/per_layer_embeddings"}, + .preshape = {config.vocab_size, + config.num_layers * config.ple_dim}, + .axes = {0, 1}, + .shape = {config.vocab_size, + config.num_layers * config.ple_dim}, + .min_size = Type::kBF16, + }); + Add(no_suffix, + { + .base_name = "ple_model_proj", + .source_names = {"embedder/per_layer_model_projection/w"}, + .preshape = {config.model_dim, + config.num_layers * config.ple_dim}, + .axes = {1, 0}, + .shape = {config.num_layers * config.ple_dim, + config.model_dim}, + .min_size = Type::kBF16, + }); + Add(no_suffix, { + .base_name = "ple_proj_norm", + .source_names = + {"embedder/per_layer_projection_norm/scale"}, + .axes = {0}, + .shape = {config.ple_dim}, + .min_size = Type::kBF16, + }); + } } // Returns the tensors for the given image layer config. @@ -407,6 +440,30 @@ void TensorInfoRegistry::AddLayerTensors(const ModelConfig& config, .shape = {1}, .min_size = Type::kBF16, }); + // Per-Layer Embedding (PLE) per-layer tensors. + if (layer_config.ple_dim > 0) { + Add(suffix, + { + .base_name = "ple_gate", + .source_names = {"per_layer_input_gate/w"}, + .axes = {1, 0}, + .shape = {layer_config.ple_dim, config.model_dim}, + }); + Add(suffix, + { + .base_name = "ple_proj", + .source_names = {"per_layer_projection/w"}, + .axes = {1, 0}, + .shape = {config.model_dim, layer_config.ple_dim}, + }); + Add(suffix, { + .base_name = "post_ple_ns", + .source_names = {"post_per_layer_input_norm/scale"}, + .axes = {0}, + .shape = {config.model_dim}, + .min_size = Type::kBF16, + }); + } Add(suffix, { .base_name = "ffw_gat_b", .source_names = {"mlp_block/ffw_up/b"}, diff --git a/gemma/tiled_attention.cc b/gemma/tiled_attention.cc index 22e20e0c..b11dff74 100644 --- a/gemma/tiled_attention.cc +++ b/gemma/tiled_attention.cc @@ -116,6 +116,8 @@ static HWY_INLINE void ComputeQKVTransposedTile( CallMatMul(activations.pre_att_rms_out, layer.qkv_einsum_w1, /*add=*/nullptr, env, activations.q); + if (flags & kSkipKV) return; + // Compute the combined KV output from pre_att_rms_out. // The output shape is [num_interleaved, kv_heads * 2 * qkv_dim]. const size_t kv_out_cols = kv_heads * 2 * qkv_dim; @@ -150,9 +152,8 @@ static HWY_INLINE void ComputeQKVTransposedTile( const size_t start_pos = qbatch.Pos(query_idx); const bool is_global_layer = activations.config.IsGlobalLayer(layer_idx); - std::vector kv_ptrs = - qbatch.KV(query_idx).cache->GetPointers( - layer_idx, kv_head, kv_heads, start_pos, is_global_layer); + std::vector kv_ptrs = qbatch.KV(query_idx).cache->GetPointers( + layer_idx, kv_head, kv_heads, start_pos, is_global_layer); const size_t v_offset = qkv_dim * KVCache::kTileSize; const size_t tile_span_size = 2 * qkv_dim * KVCache::kTileSize; const size_t k_size = qkv_dim * KVCache::kTileSize; @@ -211,10 +212,17 @@ static HWY_INLINE void ComputeQKVTransposedTile( /*mul=*/1.0f); const size_t in_tile_idx = current_pos_mod % KVCache::kTileSize; + const float* v_source = v_values; + HWY_ALIGN float v_norm_buf[kMaxQKVDim]; + if (layer_config.norm_v) { + hwy::CopyBytes(v_values, v_norm_buf, qkv_dim * sizeof(float)); + RMSNormNoScaleInplace(v_norm_buf, qkv_dim, env.ctx, worker); + v_source = v_norm_buf; + } // `v_cache_values` is a pointer to the V data that will be // compressed and stored in the KV cache. By default, it points to - // the raw `v_values`. - const float* v_cache_values = v_values; + // the raw `v_source`. + const float* v_cache_values = v_source; // `v_buf` is a temporary buffer used only when quantizing V values // to int8_t. HWY_ALIGN float v_buf[kMaxQKVDim]; @@ -249,9 +257,9 @@ static HWY_INLINE void ComputeQKVTransposedTile( // K Scaling scale_and_store(k_f32, qkv_dim, in_tile_idx); - // V Scaling: Copy `v_values` to `v_buf`, scale `v_buf` in-place, + // V Scaling: Copy `v_source` to `v_buf`, scale `v_buf` in-place, // and then update `v_cache_values` to point to `v_buf`. - hwy::CopyBytes(v_values, v_buf, qkv_dim * sizeof(float)); + hwy::CopyBytes(v_source, v_buf, qkv_dim * sizeof(float)); scale_and_store(v_buf, qkv_dim, KVCache::kTileSize + in_tile_idx); v_cache_values = v_buf; } @@ -637,9 +645,9 @@ void LocalAttentionForAllHeadsTokensAndBatch( size_t num_query_tasks = hwy::DivCeil(num_queries, kQueriesPerSubtask); [[maybe_unused]] size_t num_tasks = qbatch.Size() * layer.layer_config.kv_heads * num_query_tasks; - [[maybe_unused]] size_t num_sub_tasks = - qbatch.Size() * layer.layer_config.kv_heads * num_query_tasks * - task_multiplier; + [[maybe_unused]] size_t num_sub_tasks = qbatch.Size() * + layer.layer_config.kv_heads * + num_query_tasks * task_multiplier; HWY_DASSERT_M(activations.q.Rows() == num_query_tokens * qbatch.Size(), "qbatch size mismatch"); size_t qkv_dim = layer.layer_config.qkv_dim; @@ -663,8 +671,8 @@ void LocalAttentionForAllHeadsTokensAndBatch( if (activations.int16_queries != nullptr && num_sub_tasks * max_queries_per_subtask * qkv_dim > activations.int16_queries->size()) { - activations.int16_queries->resize(num_sub_tasks * max_queries_per_subtask * - qkv_dim); + activations.int16_queries->resize(num_sub_tasks * + max_queries_per_subtask * qkv_dim); } if (activations.q_scales != nullptr && num_sub_tasks * max_queries_per_subtask > @@ -687,8 +695,8 @@ void LocalAttentionForAllHeadsTokensAndBatch( if (activations.float_queries != nullptr && num_sub_tasks * max_queries_per_subtask * qkv_dim > activations.float_queries->size()) { - activations.float_queries->resize(num_sub_tasks * max_queries_per_subtask * - qkv_dim); + activations.float_queries->resize(num_sub_tasks * + max_queries_per_subtask * qkv_dim); } } std::vector skip_sub_task(num_sub_tasks, 0); @@ -704,10 +712,8 @@ void LocalAttentionForAllHeadsTokensAndBatch( size_t sub_task_idx = task_idx % task_multiplier; size_t query_task_idx = main_task_idx % num_query_tasks; size_t qbatch_and_kv_head_idx = main_task_idx / num_query_tasks; - size_t current_qbatch_idx = - div_kv_heads.Divide(qbatch_and_kv_head_idx); - size_t kv_head_idx = - div_kv_heads.Remainder(qbatch_and_kv_head_idx); + size_t current_qbatch_idx = div_kv_heads.Divide(qbatch_and_kv_head_idx); + size_t kv_head_idx = div_kv_heads.Remainder(qbatch_and_kv_head_idx); // First and last context token we will attend to. size_t global_start_context_pos = StartPos( qbatch.Pos(current_qbatch_idx), activations.config, layer_idx); @@ -741,7 +747,8 @@ void LocalAttentionForAllHeadsTokensAndBatch( start_context_pos + context_tokens_per_sub_task - 1); // pre-initialize memory [to avoid racy resizes laters]. size_t query_start_idx = query_task_idx * kQueriesPerSubtask; - size_t query_end_idx = std::min(num_queries, query_start_idx + kQueriesPerSubtask); + size_t query_end_idx = + std::min(num_queries, query_start_idx + kQueriesPerSubtask); size_t sub_num_queries = query_end_idx - query_start_idx; std::vector queries_ptrs; queries_ptrs.reserve(sub_num_queries); @@ -797,8 +804,7 @@ void LocalAttentionForAllHeadsTokensAndBatch( hwy::RoundDownTo(global_start_context_pos, KVCache::kTileSize); for (size_t q_idx = query_start_idx; q_idx < query_end_idx; ++q_idx) { size_t token_idx = div_heads_per_kv_head.Divide(q_idx); - int64_t global_query_pos = - qbatch.Pos(current_qbatch_idx) + token_idx; + int64_t global_query_pos = qbatch.Pos(current_qbatch_idx) + token_idx; // Intersect context to attend to for this specific query token // to the context tokens of the current subtask. int64_t query_last_context_pos = std::min( @@ -815,15 +821,16 @@ void LocalAttentionForAllHeadsTokensAndBatch( // Turn token position into KV-tile relative token positions. query_last_context_pos -= rounded_down_global_start_pos; query_start_context_pos -= rounded_down_global_start_pos; - start_pos_per_query.push_back(static_cast(query_start_context_pos)); - last_pos_per_query.push_back(static_cast(query_last_context_pos)); + start_pos_per_query.push_back( + static_cast(query_start_context_pos)); + last_pos_per_query.push_back( + static_cast(query_last_context_pos)); } if (attention_impl == AttentionImpl::kFlashTransposedQsBF16) { HWY_DASSERT(activations.bf16_queries != nullptr); - BF16* bf16_queries_ptr = - activations.bf16_queries->data() + - task_idx * max_queries_per_subtask * qkv_dim; + BF16* bf16_queries_ptr = activations.bf16_queries->data() + + task_idx * max_queries_per_subtask * qkv_dim; CompressQueriesBF16(queries_ptrs_span, qkv_dim, bf16_queries_ptr); DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsBF16( kv_ptrs, sub_num_queries, bf16_queries_ptr, @@ -838,8 +845,8 @@ void LocalAttentionForAllHeadsTokensAndBatch( int16_t* int16_queries_ptr = activations.int16_queries->data() + task_idx * max_queries_per_subtask * qkv_dim; - float* q_scales_ptr = activations.q_scales->data() + - task_idx * max_queries_per_subtask; + float* q_scales_ptr = + activations.q_scales->data() + task_idx * max_queries_per_subtask; CompressQueriesInt16(queries_ptrs_span, qkv_dim, int16_queries_ptr, q_scales_ptr); DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsInt16( @@ -851,9 +858,8 @@ void LocalAttentionForAllHeadsTokensAndBatch( max_logits.data()); } else if (attention_impl == AttentionImpl::kFlashMatrixAccumulation) { HWY_DASSERT(activations.bf16_queries != nullptr); - BF16* bf16_queries_ptr = - activations.bf16_queries->data() + - task_idx * max_queries_per_subtask * qkv_dim; + BF16* bf16_queries_ptr = activations.bf16_queries->data() + + task_idx * max_queries_per_subtask * qkv_dim; CompressAndTransposeQueriesMatrixAccumulationNonContiguous( queries_ptrs_span, bf16_queries_ptr, qkv_dim); DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsMatrixAccumulation( @@ -868,8 +874,8 @@ void LocalAttentionForAllHeadsTokensAndBatch( int8_t* int8_queries_ptr = activations.int8_queries->data() + task_idx * max_queries_per_subtask * qkv_dim; - float* q_scales_ptr = activations.q_scales->data() + - task_idx * max_queries_per_subtask; + float* q_scales_ptr = + activations.q_scales->data() + task_idx * max_queries_per_subtask; CompressAndQuantizeQueriesMatrixAccumulationInt8NonContiguous( queries_ptrs_span, int8_queries_ptr, q_scales_ptr, qkv_dim); @@ -908,13 +914,12 @@ void LocalAttentionForAllHeadsTokensAndBatch( [&](size_t main_task_idx, size_t worker) HWY_ATTR { size_t query_task_idx = main_task_idx % num_query_tasks; size_t qbatch_and_kv_head_idx = main_task_idx / num_query_tasks; - size_t current_qbatch_idx = - div_kv_heads.Divide(qbatch_and_kv_head_idx); - size_t kv_head_idx = - div_kv_heads.Remainder(qbatch_and_kv_head_idx); + size_t current_qbatch_idx = div_kv_heads.Divide(qbatch_and_kv_head_idx); + size_t kv_head_idx = div_kv_heads.Remainder(qbatch_and_kv_head_idx); size_t query_start_idx = query_task_idx * kQueriesPerSubtask; - size_t query_end_idx = std::min(num_queries, query_start_idx + kQueriesPerSubtask); + size_t query_end_idx = + std::min(num_queries, query_start_idx + kQueriesPerSubtask); for (size_t q_idx = query_start_idx; q_idx < query_end_idx; ++q_idx) { size_t sub_q_idx = q_idx - query_start_idx; diff --git a/gemma/tokenizer.cc b/gemma/tokenizer.cc index c75af1b4..b48fbf0e 100644 --- a/gemma/tokenizer.cc +++ b/gemma/tokenizer.cc @@ -111,6 +111,10 @@ GemmaChatTemplate::GemmaChatTemplate(const GemmaTokenizer& tokenizer, sot_user_ = {105, 2364, 107}; sot_model_ = {105, 4368, 107, 100, 45518, 107, 101}; eot_ = {106, 107}; + } else if (model == Model::GEMMA4_2B) { + sot_user_ = {105, 2364, 107}; + sot_model_ = {105, 4368, 107}; + eot_ = {106, 107}; } else { sot_user_.reserve(3); if (!tokenizer.Encode("user\n", &sot_user_)) return; diff --git a/gemma/vit.cc b/gemma/vit.cc index ecc0062e..8513473e 100644 --- a/gemma/vit.cc +++ b/gemma/vit.cc @@ -306,6 +306,7 @@ void FFWVit(const LayerWeightsPtrs& layer, Activations& activations, MatMulEnv& env) { PROFILER_ZONE("Gen.FFW.ViT"); const LayerConfig& layer_config = layer.layer_config; + activations.C1.OverrideCols(layer_config.ff_hidden_dim); const bool add_bias = layer_config.ff_biases; const float* bias1 = add_bias ? layer.vit.linear_0_b.PackedScale1() : nullptr; diff --git a/gemma/weights.h b/gemma/weights.h index 20932d53..ad004f4a 100644 --- a/gemma/weights.h +++ b/gemma/weights.h @@ -155,6 +155,10 @@ struct LayerWeightsPtrs { pre_ffw2_ns(finder_("pre_ffw2_ns")), moe_router(finder_("moe_router")), + ple_gate(finder_("ple_gate")), + ple_proj(finder_("ple_proj")), + post_ple_ns(finder_("post_ple_ns")), + layer_config(config) { if (layer_config.IsMoE()) { for (uint32_t i = 0; i < layer_config.NumExperts(); ++i) { @@ -227,6 +231,10 @@ struct LayerWeightsPtrs { std::vector moe_gating_einsum_w2; std::vector moe_linear_w; + MatPtr ple_gate; + MatPtr ple_proj; + MatPtr post_ple_ns; + const LayerConfig& layer_config; // Calls `func(TensorArgs)` for each tensor which is in use for the @@ -296,6 +304,12 @@ struct LayerWeightsPtrs { } } + if (layer_config.ple_dim > 0) { + func(TENSOR_ARGS(ple_gate, kMustRead)); + func(TENSOR_ARGS(ple_proj, kMustRead)); + func(TENSOR_ARGS(post_ple_ns, kMustRead)); + } + if (layer_config.ff_biases) { func(TENSOR_ARGS(ffw_gating_biases, kMustRead)); func(TENSOR_ARGS(ffw_output_biases, kMustRead)); @@ -345,6 +359,9 @@ struct WeightsPtrs { vit_img_head_bias(finder_("img_head_bias")), vit_img_head_kernel(finder_("img_head_kernel")), mm_embed_norm(finder_("mm_embed_norm")), + ple_embeddings(finder_("ple_embeddings")), + ple_model_proj(finder_("ple_model_proj")), + ple_proj_norm(finder_("ple_proj_norm")), c_layers() { c_layers.reserve(config_.layer_configs.size()); for (size_t idx = 0; idx < config_.layer_configs.size(); ++idx) { @@ -381,6 +398,10 @@ struct WeightsPtrs { MatPtr mm_embed_norm; // at least BF16. + MatPtr ple_embeddings; + MatPtr ple_model_proj; + MatPtr ple_proj_norm; + std::vector c_layers; std::vector vit_layers; @@ -403,6 +424,12 @@ struct WeightsPtrs { func(TENSOR_ARGS(embedder_input_embedding, kMustRead)); func(TENSOR_ARGS(final_norm_scale, kMustRead)); + if (config_.ple_dim > 0) { + func(TENSOR_ARGS(ple_embeddings, kMustRead)); + func(TENSOR_ARGS(ple_model_proj, kMustRead)); + func(TENSOR_ARGS(ple_proj_norm, kMustRead)); + } + if (!config_.vit_config.layer_configs.empty()) { // Vit parts. func(TENSOR_ARGS(vit_encoder_norm_bias, kMustRead)); func(TENSOR_ARGS(vit_encoder_norm_scale, kMustRead)); diff --git a/python/configs.cc b/python/configs.cc index 595ba69b..29fc85ab 100644 --- a/python/configs.cc +++ b/python/configs.cc @@ -142,6 +142,7 @@ PYBIND11_MODULE(configs, py_module) { .def_readwrite("activation", &LayerConfig::activation) .def_readwrite("post_qk", &LayerConfig::post_qk) .def_readwrite("use_qk_norm", &LayerConfig::use_qk_norm) + .def_readwrite("kv_share_layer_idx", &LayerConfig::kv_share_layer_idx) .def_readwrite("internal", &LayerConfig::internal); class_(py_module, "VitConfig") diff --git a/util/mat.h b/util/mat.h index beb9305e..04fcb42d 100644 --- a/util/mat.h +++ b/util/mat.h @@ -157,7 +157,7 @@ class MatPtr : public IFields { size_t PackedBytes() const { HWY_DASSERT_M(IsPacked(), name_.c_str()); // num_elements_ already includes the NUQ tables. - return num_elements_ * element_bytes_; + return static_cast(num_elements_) * element_bytes_; } // Works for any kind of padding and element type.