fix: fallback to dot_product for autoselected attention on CPU to prevent Tokamax NotImplementedError

[chiajunglien]

Description

This PR addresses a bug where running tasks that utilize jax.eval_shape (e.g., checkpoint conversions via to_maxtext.py) on CPU-only hardware crashes abruptly due to an unsupported attention kernel device check.

DAG Error log

Shortened Error Trace:

python
      File ".../src/maxtext/layers/attention_op.py", line 2270, in __call__
        prefill_unnormalized_output, ... = self.apply_attention( ... )
      File ".../src/maxtext/layers/attention_op.py", line 1147, in apply_attention
        out = gpu_flash_attn(query, key, value, logits_scale=1.0, mask=mask)
      File ".../site-packages/tokamax/_src/ops/attention/base.py", line 400, in __call__
        return fwd_closed( ... )
      File ".../site-packages/tokamax/_src/ops/op.py", line 197, in __call__
        raise NotImplementedError(f"Not supported on {device.device_kind}.")
    NotImplementedError: Not supported on cpu.

Root Cause

The default configuration for attention in MaxText is attention="autoselected". When autoselected falls into the non-TPU branch for the prefill / train mode, it attempts to use the GPU flash attention fallback.

• Before July 24th: The codebase utilized jax.experimental.pallas.ops.gpu.attention.mha, which gracefully allowed dummy shape tracing (jax.eval_shape) on CPU devices.
• After PR Deprecate jax.experimental.pallas GPU splash #4384 (commit e12d2fd): The codebase migrated to Tokamax's GPU flash attention (tokamax_pallas_triton). Tokamax has strict hardware validations in its tracing logic. When it detects device.device_kind == 'cpu', it immediately raises NotImplementedError: Not supported on cpu, entirely breaking abstract shape evaluation on CPU nodes.

Solution

Modified the routing logic in src/maxtext/layers/attention_op.py (apply_attention).
When attention_kernel is "autoselected" and target_hardware evaluates to "cpu", it now safely falls back to standard JAX dot-product attention (apply_attention_dot). This completely bypasses the Tokamax flash attention branch on CPU devices, preventing unexpected crashes while maintaining standard performance optimizations for TPUs and GPUs.

Tests

python3 -m maxtext.checkpoint_conversion.to_maxtext \
    model_name=gemma3-4b \
    base_output_directory=gs://runner-maxtext-logs/gemma3-4b/to_maxtext/unscanned/pre-20260725T004317 \
    use_multimodal=false \
    scan_layers=false \
    hardware=cpu skip_jax_distributed_system=True \
    checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \
    --lazy_load_tensors=False \
    --eager_load_method='transformers'

Log: https://paste.googleplex.com/6472437859483648

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.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[NuojCheng]

this change gonna silently twist all results using AOT on CPU platform

[NuojCheng]

I doubt this change can pass https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/integration/aot_identical_test.py

[NuojCheng]

somehow aot identical test does not show in https://github.com/AI-Hypercomputer/maxtext/actions/runs/30232605541/job/89875651384?pr=4617