Avoid using Row() for unaligned storage.

Balazs Racz · copybara-github · commit 858f9d138fe0 · 2025-12-18T04:13:21.000-08:00
PiperOrigin-RevId: 846158406
diff --git 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) {

Original file line number	Diff line number	Diff line change
`@@ -498,7 +498,7 @@ static void HWY_INLINE FlashAttentionTileStepAndApplySoftCap(`
`498`	`498`	`hn::Store(new_max, df4, tmp_max);`
`499`	`499`	`if constexpr (kNumQueries >= 1) {`
`500`	`500`	`const VF new_max_0 = hn::Set(df, tmp_max[0]);`
`501`		`- x_0_p0 = hn::Exp(df, hn::Sub(x_0_p0 , new_max_0));`
	`501`	`+ x_0_p0 = hn::Exp(df, hn::Sub(x_0_p0, new_max_0));`
`502`	`502`	`x_0_p1 = hn::Exp(df, hn::Sub(x_0_p1, new_max_0));`
`503`	`503`	`}`
`504`	`504`	`if constexpr (kNumQueries >= 2) {`