From 79e640a9564a7f9f2daabed50257c7b1afc7f02c Mon Sep 17 00:00:00 2001
From: Ray Smith <rays@google.com>
Date: Thu, 5 Mar 2026 07:59:16 -0800
Subject: [PATCH] Fixed tsan error.

PiperOrigin-RevId: 879069355
---
 gemma/flash_attention.cc | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/gemma/flash_attention.cc b/gemma/flash_attention.cc
index f8df7e00..0401a1f5 100644
--- a/gemma/flash_attention.cc
+++ b/gemma/flash_attention.cc
@@ -1739,7 +1739,8 @@ void ComputeFlashParams(size_t num_tokens, const size_t target_parallelism,
           size_t tasks_required = params.back().v_tile_size < k4xNFVTileSize
                                       ? k4xNFVTileSize
                                       : kVTileSize;
-          if (params.back().v_tile_size + tasks_remaining < tasks_required ||
+          if ((params.back().v_tile_size + tasks_remaining < tasks_required &&
+               params.back().v_tile_size > 0) ||
               params.back().v_tile_size == kVTileSize) {
             // We don't have enough tasks remaining to fill a tile, or the
             // current tile is full so start new tile.