From b3290e049f87c10f270b32171536e7319fa13bf4 Mon Sep 17 00:00:00 2001
From: Krzysztof Rymski <krymski@google.com>
Date: Thu, 15 Oct 2026 02:59:55 -0700
Subject: [PATCH] Internal changes

PiperOrigin-RevId: 1000003419
---
 gemma/activations.h | 10 ++++++++++
 gemma/gemma.cc      |  2 ++
 gemma/gemma.h       |  1 +
 3 files changed, 13 insertions(+)

diff --git a/gemma/activations.h b/gemma/activations.h
index 2648dadf..86b2ae27 100644
--- a/gemma/activations.h
+++ b/gemma/activations.h
@@ -435,6 +435,8 @@ struct PerToken {
   uint16_t row_idx;
 };
 static_assert(sizeof(PerToken) == 8);
+// `PerToken::expert_idx` of rows that are routed to no expert (padding).
+HWY_INLINE_VAR constexpr uint16_t kNoExpert = 0xFFFF;
 // Multiple to round up experts_per_token to avoid false sharing.
 HWY_INLINE_VAR constexpr size_t kPerTokenPerLine =
     HWY_ALIGNMENT / sizeof(PerToken);
@@ -841,6 +843,14 @@ struct Activations {
   // Filled by the embedding step; empty if unused.
   std::vector<int32_t> token_ids;
 
+  // If non-empty, nonzero for rows that only pad a prompt to the length of the
+  // others in the batch (see `ComputeEmbeddingsT`). Such rows are routed to no
+  // MoE expert, so they cost no expert FLOPs. Their outputs must be ignored.
+  std::vector<uint8_t> is_padding;
+  bool IsPadding(size_t row) const {
+    return row < is_padding.size() && is_padding[row] != 0;
+  }
+
   // Speculative decoding (DeepSeek V4 MTP): if >= 0, the compressor pass
   // copies each layer's KVCache::ds_state slice to ds_state_snapshot after
   // processing this batch row. Set to the index of the last committed token
diff --git a/gemma/gemma.cc b/gemma/gemma.cc
index 0dd2b2eb..a2dec00c 100644
--- a/gemma/gemma.cc
+++ b/gemma/gemma.cc
@@ -22,6 +22,8 @@
 #include <cstddef>
 #include <cstdint>
 #include <limits>
+#include <memory>
+#include <numeric>
 #include <optional>
 #include <vector>
 
diff --git a/gemma/gemma.h b/gemma/gemma.h
index 6a660b12..1338a2a4 100644
--- a/gemma/gemma.h
+++ b/gemma/gemma.h
@@ -18,6 +18,7 @@
 
 #include <stdio.h>
 
+#include <memory>
 #include <optional>
 #include <vector>