From 1b5d8b16af7e433289b0bb65d78e6a2078826424 Mon Sep 17 00:00:00 2001 From: Emma Lien Date: Mon, 27 Jul 2026 02:24:08 +0000 Subject: [PATCH] fix(attention): fallback to dot_product for autoselected on CPU The tokamax migration (#4384) introduced a strict hardware check that causes jax.eval_shape (used heavily during checkpoint conversion) to crash with NotImplementedError when hardware=cpu is passed. This commit intercepts target_hardware == "cpu" inside attention_op.py when attention_kernel="autoselected". It safely falls back to native JAX apply_attention_dot, ensuring that CPU shape tracing and mock evaluations complete without invoking unsupported GPU-specific flash attention logic. --- src/maxtext/layers/attention_op.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index cf076b4794..5ac6fb7da7 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -1073,6 +1073,7 @@ def apply_attention( self.attention_kernel == "dot_product" or (self.attention_kernel == "autoselected" and model_mode == MODEL_MODE_AUTOREGRESSIVE) or (self.attention_kernel == "autoselected" and length < 128) + or (self.attention_kernel == "autoselected" and target_hardware == "cpu") or (self.attention_kernel == "paged") or (self.attention_kernel in ("vllm_rpa", "vllm_batched_rpa")) ):