[BUG]: `flash_attention_gpu` fails on CPU with no fallback

Bug description

Calling flash_attention_gpu from max.nn.kernels on CPU tensors in eager mode crashes
with a cryptic Mojo compilation error. A CPU flash attention implementation exists in the
repository (I think) but is not exposed through the Python API.

It would be useful to have a cpu implementation for the torch-max-backend but also for projects like Pocket-tts which are made for cpu and not gpus :)

Root Cause

flash_attention_gpu dispatches to the mo.mha.no_cache custom op
(max/python/max/nn/kernels.py:1913),
which is a GPU-only kernel. There is no CPU fallback registered.

CPU Implementation Exists

A CPU flash attention implementation exists in Mojo (I think) at:

• max/kernels/src/nn/flash_attention.mojo
  — public entry point:

  def flash_attention[dtype, rank, ...](
      q: LayoutTensor[...],
      k_shape: IndexList[rank],
      v_shape: IndexList[rank],
      mask_shape: IndexList[mask_rank],
      output: LayoutTensor[...],
      scale: Float32,
  )
  

• max/kernels/benchmarks/nn/bench_attention_cpu.mojo
  — benchmarks the CPU kernel

Expected Behavior

flash_attention_gpu (or a new kernel flash_attention) automatically falls back to the CPU Mojo kernel when tensors are on CPU

Steps to reproduce

Reproducer

import numpy as np
from max.driver import CPU
from max.experimental import functional as F
from max.experimental.tensor import Tensor
from max.nn.kernels import MHAMaskVariant, flash_attention_gpu

batch, seq_len, num_heads, head_dim = 1, 8, 2, 16
rng = np.random.default_rng(42)
q = Tensor(rng.standard_normal((batch, seq_len, num_heads, head_dim)).astype(np.float32), device=CPU())
k = Tensor(rng.standard_normal((batch, seq_len, num_heads, head_dim)).astype(np.float32), device=CPU())
v = Tensor(rng.standard_normal((batch, seq_len, num_heads, head_dim)).astype(np.float32), device=CPU())

@F.functional
def run_attention(q: Tensor, k: Tensor, v: Tensor) -> Tensor:
    return flash_attention_gpu(q, k, v, mask_variant=MHAMaskVariant.NULL_MASK, scale=0.25)

run_attention(q, k, v)

Error

Traceback (most recent call last):
  File "/projects/open_source/torch-max-backend/.venv/lib/python3.12/site-packages/max/engine/api.py", line 513, in load
    _model = self._impl.compile_from_object(
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: Graph compilation failed:
<unknown>:46:5: error: function instantiation failed
<unknown>:66:15: note: call expansion failed with parameter value(s): (..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ...)
<unknown>:149:5: note: function instantiation failed
<unknown>:159:125: note: call expansion failed with parameter value(s): ("rank": 4, ..., ..., "local_window_size": -1, ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ...)
max/kernels/src/Mogg/MOGGKernelAPI/MOGGKernelAPI.mojo:5459:9: note: function instantiation failed
max/kernels/src/Mogg/MOGGKernelAPI/MOGGKernelAPI.mojo:5513:9: note: constraint failed: only valid on GPUs
-:1:1: error: The graph compiler could not elaborate the generated KGEN


The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/projects/open_source/torch-max-backend/trying_stuff.py", line 28, in <module>
    out = run_attention(q, k, v)
          ^^^^^^^^^^^^^^^^^^^^^^
  File "/projects/open_source/torch-max-backend/.venv/lib/python3.12/site-packages/max/experimental/functional.py", line 205, in wrapped
    with contextlib.ExitStack() as stack:
         ^^^^^^^^^^^^^^^^^^^^^^
  File "/root/.local/share/uv/python/cpython-3.12.11-linux-x86_64-gnu/lib/python3.12/contextlib.py", line 610, in __exit__
    raise exc_details[1]
  File "/root/.local/share/uv/python/cpython-3.12.11-linux-x86_64-gnu/lib/python3.12/contextlib.py", line 595, in __exit__
    if cb(*exc_details):
       ^^^^^^^^^^^^^^^^
  File "/projects/open_source/torch-max-backend/.venv/lib/python3.12/site-packages/max/experimental/realization_context.py", line 393, in __exit__
    F._run(self.realize_all())
  File "/projects/open_source/torch-max-backend/.venv/lib/python3.12/site-packages/max/experimental/functional.py", line 98, in _run
    return asyncio.run(coro)
           ^^^^^^^^^^^^^^^^^
  File "/root/.local/share/uv/python/cpython-3.12.11-linux-x86_64-gnu/lib/python3.12/asyncio/runners.py", line 195, in run
    return runner.run(main)
           ^^^^^^^^^^^^^^^^
  File "/root/.local/share/uv/python/cpython-3.12.11-linux-x86_64-gnu/lib/python3.12/asyncio/runners.py", line 118, in run
    return self._loop.run_until_complete(task)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/.local/share/uv/python/cpython-3.12.11-linux-x86_64-gnu/lib/python3.12/asyncio/base_events.py", line 691, in run_until_complete
    return future.result()
           ^^^^^^^^^^^^^^^
  File "/projects/open_source/torch-max-backend/.venv/lib/python3.12/site-packages/max/experimental/realization_context.py", line 302, in realize_all
    model = _session().load(graph)
            ^^^^^^^^^^^^^^^^^^^^^^
  File "/projects/open_source/torch-max-backend/.venv/lib/python3.12/site-packages/max/engine/api.py", line 519, in load
    raise RuntimeError(
RuntimeError: Failed to compile the model. Please file an issue, all models should be correct by construction and this error should have been caught during construction.
For more detailed failure information run with the environment variable `MODULAR_MAX_DEBUG=True`.

System information

With MAX v26.3.0.dev2026032705

Start with the CPU reproducer and inspect max/python/max/nn/kernels.py:1913, then compare its dispatch with max/kernels/src/nn/flash_attention.mojo and max/kernels/benchmarks/nn/bench_attention_cpu.mojo. Done means the reported CPU eager call no longer fails with a GPU-only compilation error and uses the available CPU path, while GPU behavior remains intact.