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

[WandLZhang]

Splash attention prefetches the decoder segment ids into SMEM. SMEM is 1 MB per core. The segmented kernel therefore can't hold the ids for much more than 524,288 tokens, and above that limit the kernel doesn't compile.

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. The SMEM block mask is (sequence / block)^2. Block 512 needs 2 MB against 1 MB of SMEM. Block 2048 clears SMEM but 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.