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

[chiajunglien]

Hi @NuojCheng, thanks for the review! I've double-checked the AOT compilation behavior, and this change works safely without twisting results.

1. AOT correctly bypasses the fallback:
Even when cross-compiling on a CPU host, AOT (via train_compile.py) uses get_topology_desc(platform="tpu") to generate mock devices. These mock devices strictly retain device.platform == "tpu". Thus, target_hardware remains "tpu", properly preserving the flash or autoselected kernels without triggering a fallback to dot_product.

2. aot_identical_test.py passes flawlessly:
I just ran the test manually on my TPU VM and can confirm the generated HLO graphs remain identical and it passes perfectly:

• tests/integration/aot_identical_test.py::AotHloIdenticalTest::test_default_hlo_match PASSED

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

If it’s missing in recent CI job logs, it’s likely because the test classes are wrapped with @pytest.mark.skip_on_tpu7x. If a GitHub Action runner happens to execute on a Trillium (tpu7x) node, Pytest naturally skips it.

Let me know if this clears up the concerns!