{"record":{"id":"5e4fffafef8f2d2a","repo":"jax-ml/jax","slug":"mosaic-gpu-does-not-yet-support-amd-rocm-devices","errorCode":null,"errorMessage":"Mosaic GPU does not yet support AMD ROCm devices. Use ``compiler_params=pltriton.CompilerParams()`` for ROCm.","messagePattern":"Mosaic GPU does not yet support AMD ROCm devices\\. Use ``compiler_params=pltriton\\.CompilerParams\\(\\)`` for ROCm\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/pallas_call.py","lineNumber":937,"sourceCode":"\n    backend: Any = None\n\n    try:\n      from jax._src.pallas.mosaic_gpu import core as mgpu_core  # pyrefly: ignore[missing-import]\n      from jax._src.pallas.mosaic_gpu import pallas_call_registration as mosaic_gpu_backend  # pyrefly: ignore[missing-import]\n    except ImportError:\n      pass\n    else:\n      if (\n          isinstance(compiler_params, mgpu_core.CompilerParams)\n          or (compiler_params is None and\n              config.jax_pallas_use_mosaic_gpu.value)\n      ):\n        backend = mosaic_gpu_backend\n\n      if backend is mosaic_gpu_backend:\n        if is_rocm:\n          raise ValueError(\n              \"Mosaic GPU does not yet support AMD ROCm devices. \"\n              \"Use ``compiler_params=pltriton.CompilerParams()`` for ROCm.\"\n          )\n\n        if ctx.primitive is pallas_call_p:\n          deprecations.warn(\n              \"jax-pallas-call-mgpu\",\n              \"Using ``pl.pallas_call`` for Mosaic GPU kernels is deprecated.\"\n              \" Support for that will be removed in a future JAX version.\"\n              \" Please migrate to ``plgpu.kernel``.\",\n              stacklevel=2,\n          )\n\n    try:\n      from jax._src.pallas.triton import core as triton_core  # pyrefly: ignore[missing-import]\n      from jax._src.pallas.triton import pallas_call_registration as triton_backend  # pyrefly: ignore[missing-import]\n    except ImportError:\n      pass","sourceCodeStart":919,"sourceCodeEnd":955,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/pallas_call.py#L919-L955","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"// before\nout = pallas_call(kernel, grid, out_shape)(x)  # mosaic backend auto-selected on ROCm\n// after\nimport jax.experimental.pallas.triton as pltriton\nout = pallas_call(kernel, grid, out_shape,\n                  compiler_params=pltriton.CompilerParams())(x)","handlingStrategy":"validation","validationCode":"import jax\nif jax.devices()[0].platform == 'gpu':\n    from jax._src.lib import xla_client\n    is_rocm = xla_client._version >= 0 and 'rocm' in jax.extend.backend.get_backend().platform_version.lower()\n    if is_rocm:\n        compiler_params = pltriton.CompilerParams()  # force Triton on ROCm","typeGuard":null,"tryCatchPattern":"try:\n    pallas_call(kernel, grid, out_shape)(x)\nexcept ValueError as e:\n    if 'ROCm' in str(e):\n        out = pallas_call(kernel, grid, out_shape,\n                          compiler_params=pltriton.CompilerParams())(x)","preventionTips":["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"],"tags":["jax","pallas","rocm","amd","gpu","backend","mosaic"],"backgroundTag":"kernel-backend-not-supported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}