diff --git a/DEVELOPERS.md b/DEVELOPERS.md index 2592673d..38e30a0f 100644 --- a/DEVELOPERS.md +++ b/DEVELOPERS.md @@ -188,6 +188,23 @@ Benchmarks such as `gemma_batch_bench` and `bench_matmul` print per-zone call counts and self-times via `PROFILER_PRINT_RESULTS()`. For performance measurements we recommend `gemma_batch_bench`. +## Attention implementations + +The attention kernel can be configured at runtime via `--attention_impl`. + +* `flash` (default) +* `flash_transposed_qs` +* `flash_transposed_qs_bf16` +* `flash_transposed_qs_int16` +* `flash_transposed_qs_int8` +* `flash_matrix_accumulation` +* `int8_matrix_accumulation` + +The KV cache type is determined from the chosen implementation unless +overridden with `--kv_cache_type`. Some implementations require a specific type +and warn if asked for another. See the `KVCache` constructor in +`gemma/kv_cache.cc` for more details. + ## Debugging At the first sign of incorrect or unexpected results, we recommend running with diff --git a/evals/benchmark.cc b/evals/benchmark.cc index 69cd644a..b142fe83 100644 --- a/evals/benchmark.cc +++ b/evals/benchmark.cc @@ -76,11 +76,12 @@ int BenchmarkCrossEntropy(GemmaEnv& env, const Path& text, size_t num_tokens = std::min(prompt.size() - pos, batch_tokens); std::vector prompt_slice(prompt.begin() + pos, prompt.begin() + pos + num_tokens); - KVCache kv_cache(gemma.Config(), gemma.Inference(), + KVCache kv_cache(gemma.Config(), gemma.Inference(), env.MutableConfig(), env.MutableEnv().ctx.allocator); float entropy = ComputeCrossEntropy(*env.GetGemma(), num_tokens, prompt_slice, kv_cache, - env.MutableEnv(), env.Verbosity()); + env.MutableEnv(), env.Verbosity(), + env.MutableConfig().attention_impl); total_entropy += entropy; LogSpeedStats(time_start, pos + num_tokens); std::string text_slice = env.StringFromTokens(prompt_slice); diff --git a/evals/benchmark_helper.cc b/evals/benchmark_helper.cc index f9a846a7..f7b74381 100644 --- a/evals/benchmark_helper.cc +++ b/evals/benchmark_helper.cc @@ -184,7 +184,7 @@ float GemmaEnv::CrossEntropy(const std::string& input) { prompt.insert(prompt.begin(), BOS_ID); return ComputeCrossEntropy(*GetGemma(), /*max_generated_tokens=*/3072, prompt, MutableKVCache(), env_, - /*verbosity=*/0) / + /*verbosity=*/0, runtime_config_.attention_impl) / static_cast(input.size()); } diff --git a/evals/cross_entropy.cc b/evals/cross_entropy.cc index 5ff4e692..5ce6ad87 100644 --- a/evals/cross_entropy.cc +++ b/evals/cross_entropy.cc @@ -98,7 +98,8 @@ HWY_EXPORT(CallSoftmax); float ComputeCrossEntropy(const Gemma& gemma, size_t max_generated_tokens, const std::vector& prompt, KVCache& kv_cache, - MatMulEnv& env, int verbosity) { + MatMulEnv& env, int verbosity, + AttentionImpl attention_impl) { const BatchStreamFunc stream_token = [](size_t, size_t, int, float) { return true; }; @@ -138,6 +139,7 @@ float ComputeCrossEntropy(const Gemma& gemma, size_t max_generated_tokens, .max_generated_tokens = max_generated_tokens - 1, .temperature = 0.0f, .verbosity = verbosity, + .attention_impl = attention_impl, .batch_stream_token = stream_token, .sample_func = sample_token, }; diff --git a/evals/cross_entropy.h b/evals/cross_entropy.h index 0a143cc9..37c234bf 100644 --- a/evals/cross_entropy.h +++ b/evals/cross_entropy.h @@ -24,9 +24,11 @@ namespace gcpp { +// Defaults match a plain RuntimeConfig, so existing callers are unaffected. float ComputeCrossEntropy(const Gemma& gemma, size_t max_generated_tokens, const std::vector& prompt, KVCache& kv_cache, - MatMulEnv& env, int verbosity); + MatMulEnv& env, int verbosity, + AttentionImpl attention_impl = AttentionImpl::kFlash); } // namespace gcpp diff --git a/evals/gemma_batch_bench.cc b/evals/gemma_batch_bench.cc index 6a4f893a..ea5e9793 100644 --- a/evals/gemma_batch_bench.cc +++ b/evals/gemma_batch_bench.cc @@ -37,8 +37,7 @@ GemmaEnv* s_env = nullptr; class GemmaBatchBench : public ::testing::Test { protected: std::vector BatchGemmaReply( - const std::vector& inputs, AttentionImpl attention_impl) { - s_env->MutableConfig().attention_impl = attention_impl; + const std::vector& inputs) { s_env->MutableConfig().temperature = 0.0f; // deterministic s_env->MutableConfig().verbosity = 2; std::vector replies; @@ -131,8 +130,7 @@ TEST_F(GemmaBatchBench, RandomQuestionsBatched) { const std::vector inputs = GenerateInputs(); // Run multiple times so that auto-tuning is closer to complete. for (size_t rep = 0; rep < 4; ++rep) { - std::vector responses = - BatchGemmaReply(inputs, AttentionImpl::kFlash); + std::vector responses = BatchGemmaReply(inputs); for (size_t i = 0; i < HWY_MIN(hwy::Unpredictable1() * 3, responses.size()); ++i) { fprintf(stderr, "Rep %zu batch answer %zu '%s'\n\n", rep, i, diff --git a/evals/run_mmlu.cc b/evals/run_mmlu.cc index c6ce9723..8cbad072 100644 --- a/evals/run_mmlu.cc +++ b/evals/run_mmlu.cc @@ -129,6 +129,7 @@ void Run(GemmaEnv& env, JsonArgs& json) { .max_generated_tokens = 30, .temperature = 0.0f, .verbosity = env.Verbosity(), + .attention_impl = env.MutableConfig().attention_impl, .stream_token = stream_token, }; env.GetGemma()->Generate(runtime_config, prompt, /*pos=*/0, diff --git a/gemma/activations.h b/gemma/activations.h index 085104b5..7ebd4143 100644 --- a/gemma/activations.h +++ b/gemma/activations.h @@ -119,6 +119,11 @@ struct AttentionActivations { att_sums( MatFactory("att_sums", batch_size, config.model_dim, allocator)), + k_tile_vec(MatFactory("k_tile_vec", batch_size * layer_config.kv_heads, + KVCache::kTileSize * max_qkv_dim, allocator)), + v_tile_vec(MatFactory("v_tile_vec", batch_size * layer_config.kv_heads, + KVCache::kTileSize * max_qkv_dim, allocator)), + inv_timescale( CreateInvTimescale(allocator, layer_config.qkv_dim, layer_config.post_qk == PostQKType::HalfRope)), @@ -225,6 +230,9 @@ struct AttentionActivationsPtrs { : config(config), flash_params(flash_params), split_flash_params(split_flash_params), + sub_task_att_out(nullptr), + sub_task_exp_denominator_sums(nullptr), + sub_task_max_logits(nullptr), bf16_queries(nullptr), int16_queries(nullptr), int8_queries(nullptr), @@ -251,6 +259,11 @@ struct AttentionActivationsPtrs { att_sums = activations.att_sums; inv_timescale = activations.inv_timescale; inv_timescale_global = activations.inv_timescale_global; + k_tile_vec = activations.k_tile_vec; + v_tile_vec = activations.v_tile_vec; + sub_task_att_out = &activations.sub_task_att_out; + sub_task_exp_denominator_sums = &activations.sub_task_exp_denominator_sums; + sub_task_max_logits = &activations.sub_task_max_logits; bf16_queries = &activations.bf16_queries; int16_queries = &activations.int16_queries; int8_queries = &activations.int8_queries; diff --git a/gemma/configs.cc b/gemma/configs.cc index ee4e4386..59da044e 100644 --- a/gemma/configs.cc +++ b/gemma/configs.cc @@ -1139,6 +1139,7 @@ Model DeduceModel(const Path& blob_path, size_t layers, int layer_types) { return Model::UNKNOWN; } +// NOTE: keep the `--attention_impl` help text in `gemma_args.h` synced constexpr std::pair kAttentionImplNameToEnum[] = { {"flash", AttentionImpl::kFlash}, {"flash_transposed_qs", AttentionImpl::kFlashTransposedQs}, @@ -1160,8 +1161,15 @@ AttentionImpl GetAttentionImpl(const std::string& impl_name) { for (const auto& [name, attention_impl] : kAttentionImplNameToEnum) { if (name == impl_name) return attention_impl; } - HWY_WARN("Unknown attention implementation: %s. Using kFlash.\n", - impl_name.c_str()); + std::string valid; + for (const auto& [name, attention_impl] : kAttentionImplNameToEnum) { + if (!valid.empty()) { + valid += ", "; + } + valid += name; + } + HWY_WARN("Unknown attention implementation: %s. Valid: %s. Using kFlash.\n", + impl_name.c_str(), valid.c_str()); return AttentionImpl::kFlash; } diff --git a/gemma/gemma_args.h b/gemma/gemma_args.h index e10bb75d..785f0b78 100644 --- a/gemma/gemma_args.h +++ b/gemma/gemma_args.h @@ -271,7 +271,11 @@ struct InferenceArgs : public ArgsBase { "When a newline is encountered, that signals the end of the turn.", 2); visitor(attention_impl, "attention_impl", std::string("flash"), - "Attention implementation to use. See configs.cc for options.", 2); + "Attention implementation (flash, flash_transposed_qs, " + "flash_transposed_qs_bf16, flash_transposed_qs_int16, " + "flash_transposed_qs_int8, flash_matrix_accumulation, " + "int8_matrix_accumulation).", + 2); visitor(kv_cache_type, "kv_cache_type", std::string(""), "KV cache data type (f32, bf16, int8). If empty, deduced from " "attention_impl.",