Let splash attention exceed 524,288 tokens when packing is off

[WandLZhang]

At max_target_length=1048576 the segmented splash kernel fails to compile:

RESOURCE_EXHAUSTED: Allocation (size=2097152) would exceed memory (size=1048576)
:: #allocation3 [shape = 'u8[2097152]{0}', space=smem, size = 0x200000, scoped,
   tag = 'prefetched SMEM operand 0'] :: splash_mha_fwd_segmented_residuals.3

I can't say which array that is. make_splash_mha uses num_scalar_prefetch = 3 and passes fwd_mask_info.data_next, block_mask and mask_next as the first three operands, so operand 0 is data_next. But 2,097,152 has several exact readings: data_next as int16[1, 1024, 1024] at block 1024, or a two-byte segment-id array at 1M tokens. I've stopped deriving it from the number.

Setting decoder_segment_ids to None when packing is off clears the failure. It's also correct on its own terms: with packing off there's one segment per example, the ids are constant, and that's what the non-segmented kernel assumes.

Four later runs on 64 v5p chips at 1M cross segment ids against splash block size. All four train, so I can't reproduce the failure above on that machine:

segment ids	splash block	tokens/s/chip
dropped	1024	513
present	1024	512
dropped	512	423
present	512	412

Treat the allocation dump as the report and the mechanism as open. Block 1024 is worth 21% at 1M.

When packing is off, each example is a single segment. The segment ids are constant, so the non-segmented kernel produces the same result and has no such limit. This change sets decoder_segment_ids_q to None when config.packing is false.

Effect

Measured on a v5p with a Qwen 3.5 hybrid model and context parallelism:

Sequence	Before	After
524,288	compiles	compiles
1,048,576	kernel doesn't compile	trains end to end

Set the splash block sizes to 1024 above 524,288 tokens. On this config it is worth 21% at 1M, 513 against 423 tokens/s/chip. Block 2048 exceeds VMEM.

Scope

The change applies only when packing is false. Packed runs keep the segmented kernel and are unaffected.

Tests

No unit test covers this path. I didn't add one, because the test needs a TPU and a sequence long enough to fill SMEM. Let me know where such a test should live and I'll write it.

cc @mmcsa

[gemini-code-assist]

Code Review

This pull request drops decoder segment IDs when packing is disabled to avoid a token limit in Splash attention. However, the reviewer noted that the check is placed too late in wrap_flash_attention, bypassing several early return paths. It is recommended to move this check to the beginning of the function and apply it to both query and key-value segment IDs.

[gemini-code-assist]

The current placement of the if not self.config.packing: check is too late because it occurs after the early returns for use_tokamax_ring (line 1999), use_ulysses (line 2025), and use_usp (line 2036). As a result, the segment IDs are not dropped for these attention paths, and they will still hit the 524,288 token limit when packing is off.

To fix this, please move this check and its accompanying comment to the very beginning of wrap_flash_attention (right after line 1988). Additionally, you should set both decoder_segment_ids_q and decoder_segment_ids_kv to None to ensure consistency across all paths.

Here is how it should look at the top of wrap_flash_attention:

    def wrap_flash_attention(
        query,
        key,
        value,
        decoder_segment_ids_q,
        decoder_segment_ids_kv,
        sa_config,
        splash_kernel,
        cp_size,
        load_balanced_context_parallel,
        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 in that case to remove the limit.
      if not self.config.packing:
        decoder_segment_ids_q = None
        decoder_segment_ids_kv = None

[WandLZhang]

Moved the check to the top of wrap_flash_attention, before the use_tokamax_ring, use_ulysses and use_usp returns, and cleared decoder_segment_ids_kv as well as decoder_segment_ids_q.

[WandLZhang]

Converting to draft. The failing config was 16 chips at ctx=16; the four runs I used to probe it were 64 chips at ctx=64, so they never reproduced the failure and I can't state the mechanism beyond what the error says — operand 0 is data_next, 2 MB against 1 MB of SMEM.

The block-1024 throughput result stands on its own (513 against 423 tokens/s/chip at 1M) and I'll re-raise this separately if I can reproduce the failure.