From 0ac55f71edb7a190197e18c5deb22d921e0e5821 Mon Sep 17 00:00:00 2001
From: Balazs Racz <bracz@google.com>
Date: Thu, 18 Dec 2025 05:10:21 -0800
Subject: [PATCH] Avoid using Row() for unaligned storage.

PiperOrigin-RevId: 846214605
---
 gemma/flash_attention.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/gemma/flash_attention.cc b/gemma/flash_attention.cc
index 06e000a2..c56c036f 100644
--- a/gemma/flash_attention.cc
+++ b/gemma/flash_attention.cc
@@ -498,7 +498,7 @@ static void HWY_INLINE FlashAttentionTileStepAndApplySoftCap(
   hn::Store(new_max, df4, tmp_max);
   if constexpr (kNumQueries >= 1) {
     const VF new_max_0 = hn::Set(df, tmp_max[0]);
-    x_0_p0 = hn::Exp(df, hn::Sub(x_0_p0 , new_max_0));
+    x_0_p0 = hn::Exp(df, hn::Sub(x_0_p0, new_max_0));
     x_0_p1 = hn::Exp(df, hn::Sub(x_0_p1, new_max_0));
   }
   if constexpr (kNumQueries >= 2) {