diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py
index 81beb97020..e0d9410142 100644
--- a/src/maxtext/layers/attention_op.py
+++ b/src/maxtext/layers/attention_op.py
@@ -1986,6 +1986,15 @@ def wrap_flash_attention(
         sinks,
         indexer_mask,
     ):
+      # Splash prefetches the segment ids into SMEM, and SMEM is 1 MB per core.
+      # This limits the segmented kernel to approximately 524,288 tokens. When
+      # packing is off there is one segment per example, so the segment ids are
+      # constant and the non-segmented kernel gives the same result. Drop the
+      # ids here, before any early return, so every path benefits.
+      if not self.config.packing:
+        decoder_segment_ids_q = None
+        decoder_segment_ids_kv = None
+
       if use_tokamax_ring:
         attention_output = tokamax_ring_attention.call_ring_attention(
             query,