From 43f4bcef0ed1452b3c339671aeb847c3fed850ed Mon Sep 17 00:00:00 2001 From: Willis Zhang Date: Tue, 18 Aug 2026 09:36:45 -0400 Subject: [PATCH 1/2] Let splash attention exceed 524,288 tokens when packing is off Splash prefetches the segment ids into SMEM. SMEM is 1 MB per core, so the segmented kernel cannot hold the ids for much more than 524,288 tokens. Above that the kernel fails to compile. When packing is off, each example is one segment. The segment ids are therefore constant, and the non-segmented kernel produces the same result without the SMEM cost. This change drops the ids in that case. Measured on a v5p: with this change a sequence of 1,048,576 tokens compiles and trains. Without it the same configuration fails at 524,288. --- src/maxtext/layers/attention_op.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index 81beb97020..c2b3e20526 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -2056,6 +2056,14 @@ def wrap_flash_attention( to_contiguous=True, ) + # 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 in that case to remove the limit. + if not self.config.packing: + decoder_segment_ids_q = None + if decoder_segment_ids_q is not None: if cp_size > 1 and load_balanced_context_parallel: decoder_segment_ids_tuple = splash_attention_kernel.SegmentIds( From 83f18aab952d23d8f291153c9541736a74c8edd0 Mon Sep 17 00:00:00 2001 From: Willis Zhang Date: Tue, 18 Aug 2026 09:58:59 -0400 Subject: [PATCH 2/2] Drop the segment ids before the early returns The check sat after the early returns for tokamax_ring, ulysses and usp, so those three paths kept the segment ids and still hit the SMEM limit. Move it to the top of wrap_flash_attention, and clear decoder_segment_ids_kv as well as decoder_segment_ids_q so every path is consistent. Caught in review by gemini-code-assist on #4929. --- src/maxtext/layers/attention_op.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index c2b3e20526..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, @@ -2056,14 +2065,6 @@ def wrap_flash_attention( to_contiguous=True, ) - # 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 in that case to remove the limit. - if not self.config.packing: - decoder_segment_ids_q = None - if decoder_segment_ids_q is not None: if cp_size > 1 and load_balanced_context_parallel: decoder_segment_ids_tuple = splash_attention_kernel.SegmentIds(