jax-ml/jax · error · ValueError

Mosaic GPU does not yet support AMD ROCm devices. Use ``comp

Error message

Mosaic GPU does not yet support AMD ROCm devices. Use ``compiler_params=pltriton.CompilerParams()`` for ROCm.

What it means

The Mosaic GPU compiler backend does not support AMD ROCm devices. If the Mosaic GPU backend is selected and the current GPU is ROCm, pallas_call lowering raises ValueError directing the user to the Triton-based compiler backend for ROCm support.

Source

Thrown at jax/_src/pallas/pallas_call.py:937

    backend: Any = None

    try:
      from jax._src.pallas.mosaic_gpu import core as mgpu_core  # pyrefly: ignore[missing-import]
      from jax._src.pallas.mosaic_gpu import pallas_call_registration as mosaic_gpu_backend  # pyrefly: ignore[missing-import]
    except ImportError:
      pass
    else:
      if (
          isinstance(compiler_params, mgpu_core.CompilerParams)
          or (compiler_params is None and
              config.jax_pallas_use_mosaic_gpu.value)
      ):
        backend = mosaic_gpu_backend

      if backend is mosaic_gpu_backend:
        if is_rocm:
          raise ValueError(
              "Mosaic GPU does not yet support AMD ROCm devices. "
              "Use ``compiler_params=pltriton.CompilerParams()`` for ROCm."
          )

        if ctx.primitive is pallas_call_p:
          deprecations.warn(
              "jax-pallas-call-mgpu",
              "Using ``pl.pallas_call`` for Mosaic GPU kernels is deprecated."
              " Support for that will be removed in a future JAX version."
              " Please migrate to ``plgpu.kernel``.",
              stacklevel=2,
          )

    try:
      from jax._src.pallas.triton import core as triton_core  # pyrefly: ignore[missing-import]
      from jax._src.pallas.triton import pallas_call_registration as triton_backend  # pyrefly: ignore[missing-import]
    except ImportError:
      pass

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Pass compiler_params=jax.experimental.pallas.triton.CompilerParams() so the Triton backend compiles for ROCm
Do not enable the Mosaic GPU backend (unset jax_pallas_use_mosaic_gpu / avoid mosaic-specific compiler params) on AMD hardware
Run the workload on a CUDA (NVIDIA) or TPU device if Mosaic is required
Update JAX — ROCm/Mosaic support status changes between versions

Example fix

// before
out = pallas_call(kernel, grid, out_shape)(x)  # mosaic backend auto-selected on ROCm
// after
import jax.experimental.pallas.triton as pltriton
out = pallas_call(kernel, grid, out_shape,
                  compiler_params=pltriton.CompilerParams())(x)

Defensive patterns

Strategy: validation

Validate before calling

import jax
if jax.devices()[0].platform == 'gpu':
    from jax._src.lib import xla_client
    is_rocm = xla_client._version >= 0 and 'rocm' in jax.extend.backend.get_backend().platform_version.lower()
    if is_rocm:
        compiler_params = pltriton.CompilerParams()  # force Triton on ROCm

Try / catch

try:
    pallas_call(kernel, grid, out_shape)(x)
except ValueError as e:
    if 'ROCm' in str(e):
        out = pallas_call(kernel, grid, out_shape,
                          compiler_params=pltriton.CompilerParams())(x)

Prevention

Detect ROCm and set Triton CompilerParams in cluster init code
Do not enable jax_pallas_use_mosaic_gpu globally on heterogeneous GPU clusters
Pin backend-specific kernel configs per hardware target

When it happens

Trigger: Running a pallas_call that selects the mosaic GPU backend (or has jax_pallas_use_mosaic_gpu enabled) on a machine with an AMD GPU using ROCm; e.g. importing a Mosaic-targeted kernel on an MI2xx/MI3xx system.

Common situations: Running TPU-oriented Mosaic kernels or libraries on AMD GPU clusters; enabling the Mosaic GPU experimental flag globally and then running the same code on a heterogeneous cluster with NVIDIA and AMD nodes.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/5e4fffafef8f2d2a. Report an issue: GitHub.