Internal changes

BUILD.bazel

@@ -780,7 +780,6 @@ cc_library(
         ":ops",
         ":tensor_stats",
         ":threading_context",
-        "@highway//:abort_header_only",
     ],
 )
 
@@ -924,6 +923,7 @@ cc_library(
         ":query",
         ":tensor_stats",
         ":threading_context",
+        ":tokenizer",
         ":weights",
         ":zones",
         "//compression:compress",
@@ -1031,7 +1031,9 @@ cc_library(
 
 cc_test(
     name = "gemma_test",
-    srcs = ["evals/gemma_test.cc"],
+    srcs = [
+        "evals/gemma_test.cc",
+    ],
     linkstatic = True,
     # Requires model files
     tags = [

evals/gemma_test.cc

@@ -17,6 +17,7 @@
 
 #include <stdio.h>
 
+#include <cmath>
 #include <string>
 #include <vector>
 
@@ -41,8 +42,17 @@ class GemmaTest : public ::testing::Test {
   // Requires argc/argv, hence do not use `SetUpTestSuite`.
   static void InitEnv(int argc, char** argv) {
     HWY_ASSERT(s_env == nullptr);  // Should only be called once.
-    ConsumedArgs consumed(argc, argv);
-    GemmaArgs args(argc, argv, consumed);
+    std::vector<char*> filtered_argv;
+    for (int i = 0; i < argc; ++i) {
+      // Test runners may pass `--logtostderr`, which is not
+      // recognized by `GemmaArgs`. Filter it out so
+      // `ConsumedArgs::AbortIfUnconsumed` does not abort.
+      if (std::string_view(argv[i]) == "--logtostderr") continue;
+      filtered_argv.push_back(argv[i]);
+    }
+    int filtered_argc = static_cast<int>(filtered_argv.size());
+    ConsumedArgs consumed(filtered_argc, filtered_argv.data());
+    GemmaArgs args(filtered_argc, filtered_argv.data(), consumed);
     consumed.AbortIfUnconsumed();
 
     s_env = new GemmaEnv(args);
@@ -75,6 +85,9 @@ class GemmaTest : public ::testing::Test {
 GemmaEnv* GemmaTest::s_env = nullptr;
 
 TEST_F(GemmaTest, Batched) {
+  if (s_env->GetGemma()->Config().IsEmbedding()) {
+    GTEST_SKIP() << "Not applicable for embedding models";
+  }
   // Test remainder handling in MatMul (four rows per tile), but avoid a
   // second batch in debug builds to speed up the test.
   s_env->MutableConfig().decode_qbatch_size = HWY_IS_DEBUG_BUILD ? 6 : 3;
@@ -102,6 +115,9 @@ TEST_F(GemmaTest, Batched) {
 TEST_F(GemmaTest, Multiturn) {
   const Gemma* model = s_env->GetGemma();
   const ModelConfig& config = model->Config();
+  if (config.IsEmbedding()) {
+    GTEST_SKIP() << "Not applicable for embedding models";
+  }
   size_t abs_pos = 0;
   std::string response;
   auto stream_token = [&](size_t query_idx, size_t pos, int token, float) {
@@ -156,6 +172,9 @@ TEST_F(GemmaTest, Multiturn) {
 TEST_F(GemmaTest, CrossEntropySmall) {
   HWY_ASSERT(s_env->GetGemma() != nullptr);
   const ModelConfig& config = s_env->GetGemma()->Config();
+  if (config.IsEmbedding()) {
+    GTEST_SKIP() << "Not applicable for embedding models";
+  }
   static const char kSmall[] =
       "The capital of Hungary is Budapest which is located in Europe.";
   float entropy = s_env->CrossEntropy(kSmall);

gemma/configs.h

@@ -810,6 +810,7 @@ struct ModelConfig : public IFields {
   }
 
   bool IsEOS(int id) const { return (id == eos_id || id == secondary_eos_id); }
+  bool IsEmbedding() const { return false; }
 
   // Major version of the model family, reflecting architecture changes. This is
   // more convenient to compare than `Model` because that also includes the

gemma/flash_attention.cc

@@ -2146,7 +2146,9 @@ void ComputeFlashParams(size_t num_tokens, const size_t target_parallelism,
         const size_t prefix_end = qbatch.PrefixEnd(qi);
         if (prefix_end > 0 && prefix_end - 1 > last) {
           // last_pos is inclusive.
-          last = prefix_end - 1;
+          const size_t window_size =
+              activations.config.attention_window_sizes[layer_idx];
+          last = HWY_MIN(prefix_end - 1, pos + window_size - 1);
         }
         for (size_t head_group = 0; head_group < kHeadGroups; ++head_group) {
           size_t tasks_remaining = kHeadGroups - head_group +

gemma/gemma.h

@@ -35,6 +35,7 @@
 #include "paligemma/image.h"
 #include "util/basics.h"  // TokenAndProb
 #include "util/threading_context.h"
+#include "hwy/aligned_allocator.h"  // AlignedVector
 #include "hwy/timer.h"
 // IWYU pragma: end_exports
 

gemma/tiled_attention.cc

@@ -213,6 +213,9 @@ static HWY_INLINE void ComputeQKVTransposedTile(
                 RMSNormInplace(weights_t->PackedScale1(), /*w_ofs=*/0, k_f32,
                                qkv_dim, env.ctx, worker);
               });
+            } else if (layer_config.post_qk == PostQKType::NormLocalRope ||
+                       layer_config.use_qk_norm) {
+              RMSNormNoScaleInplace(k_f32, qkv_dim, env.ctx, worker);
             }
             PositionalEncodingQK(
                 k_f32, layer_idx, activations, env.ctx, worker,
@@ -930,7 +933,10 @@ void LocalAttentionForAllHeadsTokensAndBatch(
         // that into account.
         const size_t prefix_end = qbatch.PrefixEnd(current_qbatch_idx);
         if (prefix_end > 0 && prefix_end - 1 > last_context_pos) {
-          last_context_pos = prefix_end - 1;
+          const size_t window_size =
+              activations.config.attention_window_sizes[layer_idx];
+          last_context_pos =
+              std::min(prefix_end - 1, last_context_pos + window_size - 1);
         }
         size_t total_num_context_tokens =
             last_context_pos - start_context_pos + 1;
@@ -1007,26 +1013,53 @@ void LocalAttentionForAllHeadsTokensAndBatch(
         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;
-          // 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(
-              static_cast<int64_t>(last_context_pos), global_query_pos);
-          // This max is to not go into negative values, for the same reason we
-          // use int64_t and not size_t here.
+          // Compute the range of context tokens [query_start_context_pos,
+          // query_last_context_pos] that this query token should attend to
+          // within the current KV tile/subtask.
+
+          // For standard causal attention, a token cannot attend to future
+          // positions (query_last_pos <= global_query_pos). For bidirectional
+          // prefix attention (prefix_end > 0), tokens within the prefix can
+          // attend forward to subsequent prefix tokens, capped by the local
+          // sliding window size.
+          int64_t query_last_pos = global_query_pos;
+          if (prefix_end > 0 &&
+              prefix_end - 1 > static_cast<size_t>(query_last_pos)) {
+            const size_t window_size =
+                activations.config.attention_window_sizes[layer_idx];
+            query_last_pos = std::min(
+                static_cast<int64_t>(prefix_end - 1),
+                global_query_pos + static_cast<int64_t>(window_size) - 1);
+          }
+          int64_t query_last_context_pos =
+              std::min(static_cast<int64_t>(last_context_pos), query_last_pos);
+
+          // The query cannot attend backward beyond the sliding window
+          // (global_query_pos - window_size + 1). Clamp to start_context_pos of
+          // the current subtask. Signed int64_t is used to avoid underflow when
+          // global_query_pos < window_size.
           int64_t query_start_context_pos = std::max(
               global_query_pos -
                   static_cast<int64_t>(
                       activations.config.attention_window_sizes[layer_idx]) +
                   1,
               static_cast<int64_t>(start_context_pos));
 
-          // 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<size_t>(query_start_context_pos));
-          last_pos_per_query.push_back(
-              static_cast<size_t>(query_last_context_pos));
+          // If the query's attention window does not overlap with this KV tile,
+          // set start_pos > last_pos (SIZE_MAX and 0) so the attention kernel
+          // skips this query.
+          if (query_last_context_pos < query_start_context_pos) {
+            start_pos_per_query.push_back(std::numeric_limits<size_t>::max());
+            last_pos_per_query.push_back(0);
+          } else {
+            // 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<size_t>(query_start_context_pos));
+            last_pos_per_query.push_back(
+                static_cast<size_t>(query_last_context_pos));
+          }
         }
 
         if (attention_impl == AttentionImpl::kFlashTransposedQsBF16) {