Fix packed-sequence masking in cudnn_flash_jax and add cutlass_flash GPU backend

[bzantium]

Description

This PR fixes a silent correctness bug in the cudnn_flash_jax attention kernel and adds cutlass_flash, an optional FlashAttention-2 (CUTLASS, via flash-attn-jax) backend for GPUs, together with a pallas kernel used for the shapes FA2 cannot run. The two changes are coupled: the fused fallback that makes the cudnn_flash_jax fix practical on pre-Hopper GPUs is provided by the new backend.

Problem 1: cudnn_flash_jax ignores decoder_segment_ids

AttentionOp.cudnn_jax_flash_attention received decoder_segment_ids but its train branch called the jax cuDNN SDPA with mask_type=MaskType.CAUSAL only. With sequence packing enabled (the default packing: true with any real dataset), tokens attended across packed segment boundaries: no guard, no error, silently wrong loss and gradients. In a kernel-vs-reference test on packed batches, every token after the packing boundary diverged from the dot_product reference (max abs diff 5.4 vs 0.009 bf16 noise on the first segment).

Fix, in order of preference at runtime:

• On Hopper and newer, lower segment ids to the cuDNN packed layout: per-segment q_seqlen/kv_seqlen [B, M] plus q_offsets/kv_offsets [B, M+1] (absent segments filled with -1) with MaskType.PADDING_CAUSAL. This is the layout jax._src.cudnn.fused_attention_stablehlo.dot_product_attention documents for packed inputs.
• On older GPUs the cuDNN packed layout is rejected by jax (Packed layout requires a GPU with at least Hopper architecture), so the op falls back to a fused flash kernel that masks segment boundaries (see below) and logs the downgrade, instead of silently attending across boundaries.
• max_segments_per_seq is now validated at config time for cudnn_flash_jax (and cutlass_flash) with packing, matching the existing cudnn_flash_te validation.

Problem 2: no fast, correct packed attention on pre-Hopper GPUs

On A100 we found no working fused path for packed sequences: the jax cuDNN packed layout is Hopper-only, and cudnn_flash_te (TE 2.17, cuDNN 9.24) finds no sm80 THD kernel and then crashes in its unfused fallback (sequence_descriptor must be None or ndarray). The remaining correct option, unfused dot_product, costs about 39 percent of step throughput at seq 2048 on a qwen3-4b benchmark.

Key changes

• attention=cutlass_flash: FlashAttention-2 CUTLASS kernels through the optional flash-attn-jax package (import-guarded; the option errors clearly when the package is absent). Native GQA/MQA, sliding window (LOCAL_SLIDING), and sequence packing lowered to FA2's varlen kernel with cu_seqlens built inside jit. Requires bf16/fp16 and head_dim <= 256; decode falls back to unfused attention.
• src/maxtext/kernels/attention/gpu_pallas_flash.py: a fork of jax.experimental.pallas.ops.gpu.attention extended with sliding-window masking, gemma2-style logit soft cap (with the matching backward chain rule), and chunked head_dim accumulation for head_dim > 256, so every layer FA2 cannot run still has a fused kernel.
• A single backend dispatch (AttentionOp._gpu_flash_backend_dispatch) now serves attention=flash (GPU), attention=cutlass_flash, and the pre-Hopper packed fallback of cudnn_flash_jax: FA2 where eligible, pallas otherwise. Eligibility covers dtype (bf16/fp16), head_dim, soft cap, and packed-metadata availability, so no configuration silently loses segment masking.
• Context parallelism is explicitly rejected for these two backends (the kernels compute masks from shard-local positions); cudnn_flash_te remains the kernel to use with context parallelism.

Measurements (A100-SXM4-80GB x 8, qwen3-4b, bf16, seq 2048, per-device batch 8, synthetic data)

attention	TFLOP/s/device	notes
cudnn_flash_te	177.8	crashes with packing on sm80 (TE fallback bug)
cutlass_flash	176.6	fast and correct with packing
dot_product	108.0	previous only correct packed option on sm80

Shortcomings / future work

• The Hopper-native packed path of cudnn_flash_jax (seqlens/offsets) is implemented per the jax API contract and covered by the new test, but our verification hardware is A100; a run of tests/unit/gpu_attention_packing_test.py on H100 or newer would exercise it directly.
• The TE unfused-fallback crash (SequenceDescriptor assertion) is a separate upstream issue and is not addressed here.

FIXES: #4476

Tests

New test: tests/unit/gpu_attention_packing_test.py (marked gpu_only). For each kernel (cudnn_flash_jax, cutlass_flash, and the pallas backend forced by hiding the cutlass library), it compares the forward output and the gradient wrt query against the dot_product reference on an unpacked control, a packed batch, and a packed batch with a tile-unaligned boundary, reporting relative error separately before and after the packing boundary so cross-segment leaks are directly visible. A host-side test covers the segment-ids-to-seqlens/offsets encoding.

Results on A100 (seq 2048, GQA 32/8 heads, bf16): all kernels match the reference within 2 percent relative error on both forward and gradients, for all three boundary cases. Reproduce with:

PYTHONPATH=src:. python -m pytest tests/unit/gpu_attention_packing_test.py -v

pyink (23.10.0, repo settings) and pylint (repo pylintrc) pass on the changed files.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.