From dc3914f4f607ff52283ab9c20aa2b110e16788af Mon Sep 17 00:00:00 2001 From: Krzysztof Rymski Date: Tue, 8 Sep 2026 01:33:59 -0700 Subject: [PATCH] Internal changes PiperOrigin-RevId: 977763909 --- BUILD.bazel | 6 ++-- evals/gemma_test.cc | 23 +++++++++++++-- gemma/configs.h | 1 + gemma/flash_attention.cc | 4 ++- gemma/gemma.h | 1 + gemma/tiled_attention.cc | 61 +++++++++++++++++++++++++++++++--------- 6 files changed, 77 insertions(+), 19 deletions(-) diff --git a/BUILD.bazel b/BUILD.bazel index 60bc61d6..4f58cfe1 100644 --- a/BUILD.bazel +++ b/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 = [ diff --git a/evals/gemma_test.cc b/evals/gemma_test.cc index 17ff78a0..21c369c9 100644 --- a/evals/gemma_test.cc +++ b/evals/gemma_test.cc @@ -17,6 +17,7 @@ #include +#include #include #include @@ -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 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(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); diff --git a/gemma/configs.h b/gemma/configs.h index 0a9c4078..d3222ad6 100644 --- a/gemma/configs.h +++ b/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 diff --git a/gemma/flash_attention.cc b/gemma/flash_attention.cc index 66514699..b801b96d 100644 --- a/gemma/flash_attention.cc +++ b/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 + diff --git a/gemma/gemma.h b/gemma/gemma.h index f04e8b78..6a660b12 100644 --- a/gemma/gemma.h +++ b/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 diff --git a/gemma/tiled_attention.cc b/gemma/tiled_attention.cc index f53ebb6f..a42ac29a 100644 --- a/gemma/tiled_attention.cc +++ b/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,12 +1013,31 @@ 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(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(query_last_pos)) { + const size_t window_size = + activations.config.attention_window_sizes[layer_idx]; + query_last_pos = std::min( + static_cast(prefix_end - 1), + global_query_pos + static_cast(window_size) - 1); + } + int64_t query_last_context_pos = + std::min(static_cast(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( @@ -1020,13 +1045,21 @@ void LocalAttentionForAllHeadsTokensAndBatch( 1, static_cast(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(query_start_context_pos)); - last_pos_per_query.push_back( - static_cast(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::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(query_start_context_pos)); + last_pos_per_query.push_back( + static_cast(query_last_context_pos)); + } } if (attention_impl == AttentionImpl::kFlashTransposedQsBF16) {