From 6661d3a60c3d13797f82c2bd031119cbb3fb4302 Mon Sep 17 00:00:00 2001
From: Krzysztof Rymski <krymski@google.com>
Date: Thu, 18 Dec 2025 01:26:09 -0800
Subject: [PATCH] Internal changes

PiperOrigin-RevId: 846140314
---
 BUILD.bazel              | 9 ---------
 gemma/flash_attention.cc | 4 ++--
 2 files changed, 2 insertions(+), 11 deletions(-)

diff --git a/BUILD.bazel b/BUILD.bazel
index 8c398e73..39b9fe9f 100644
--- a/BUILD.bazel
+++ b/BUILD.bazel
@@ -606,11 +606,7 @@ cc_library(
     ],
     deps = [
         ":activations",
-        ":allocator",
-        ":basics",
         ":configs",
-        ":gemma_args",
-        ":kv_cache",
         ":mat",
         ":matmul",
         ":matmul_env",
@@ -623,13 +619,8 @@ cc_library(
         ":zones",
         "//compression:compress",
         "//compression:types",
-        "//io",
-        "@highway//:bit_set",
         "@highway//:hwy",
-        "@highway//:nanobenchmark",  # timer
         "@highway//:profiler",
-        "@highway//:thread_pool",
-        "@highway//hwy/contrib/sort:vqsort",
     ],
 )
 
diff --git a/gemma/flash_attention.cc b/gemma/flash_attention.cc
index f5c5cca3..06e000a2 100644
--- a/gemma/flash_attention.cc
+++ b/gemma/flash_attention.cc
@@ -483,7 +483,7 @@ static void HWY_INLINE FlashAttentionTileStepAndApplySoftCap(
     new_max = hn::InsertLane(new_max, 0, hn::ReduceMax(df, max_0));
   } else {
     new_max = Reduce4(df, max_0, max_1, max_2, max_3,
-                             [](auto a, auto b) { return hn::Max(a, b); });
+                      [](auto a, auto b) HWY_ATTR { return hn::Max(a, b); });
   }
   if (att_cap > 0.0f) {
     VF4 cap = hn::Set(df4, att_cap);
@@ -531,7 +531,7 @@ static void HWY_INLINE FlashAttentionTileStepAndApplySoftCap(
     VF x_2_sum = hn::Add(x_2_p0, x_2_p1);
     VF x_3_sum = hn::Add(x_3_p0, x_3_p1);
     x_sum = Reduce4(df, x_0_sum, x_1_sum, x_2_sum, x_3_sum,
-                           [](auto a, auto b) { return hn::Add(a, b); });
+                    [](auto a, auto b) HWY_ATTR { return hn::Add(a, b); });
   }
   old_d_vf = hn::Add(scale, x_sum);
   auto non_zero_mask = hn::Gt(old_d_vf, hn::Set(df4, 0.0f));