jax-ml/jax · error · NotImplementedError

inline_mgpu_p does not support discharge.

Error message

inline_mgpu_p does not support discharge.

What it means

inline_mgpu primitives perform raw MLIR operations inside a Pallas kernel and have no pure-JAX fallback; attempting to discharge (lower a Pallas kernel to plain JAX, e.g. for CPU interpretation or pipelining outside GPU) is explicitly unsupported.

Source

Thrown at jax/_src/pallas/mosaic_gpu/primitives.py:3519

      jax_core.ShapedArray(x.shape, x.dtype) for x in flat_ret_ty
  )
  # TODO(cperivol): Let the user set the effects.
  flat_args = flat_args_and_transforms[:pytree_args.num_leaves]
  return aval_return, {
      gpu_core._wgmma_pipeline_effect,
      gpu_core._memory_effect,
      *itertools.chain.from_iterable(
          (state.ReadEffect(i), state.WriteEffect(i))
          for i, r in enumerate(flat_args)
          if isinstance(r, state.AbstractRef)
      ),
  }


@discharge.register_discharge_rule(inline_mgpu_p)
def _inline_mgpu_discharge(*args, **kwargs):
  del args, kwargs
  raise NotImplementedError("inline_mgpu_p does not support discharge.")


def _type_check_mgpu_lane_semantics(v, ty):
  match (ty, v):
    case (RefType(), ir.Value()) if isinstance(v.type, ir.MemRefType):
      pass
    case (ShapeDtypeStruct(), mgpu.FragmentedArray()):
      mlir_dtype = mgpu_utils.dtype_to_ir_type(ty.dtype)
      if v.mlir_dtype != mlir_dtype:
        raise ValueError(
            f"Array dtype mismatch: expected {v.mlir_dtype} got {mlir_dtype}."
        )
      if ty.shape != v.shape:
        raise ValueError(
            f"Array shape mismatch: expected {ty.shape} got {v.shape}."
        )
      if v.layout != ty.layout.to_mgpu():
        raise ValueError(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Run the kernel on GPU (compiled path) instead of the interpreter/discharge path
2. Replace inline_mgpu usage with standard Pallas/Mosaic primitives that support discharge when interpretation is required

Example fix

# before
kernel_with_inline_mgpu(...)  # interpreted/discharge run
# after
compiled_gpu_kernel(...)  # run via pallas_call with a GPU compilation target

try:
    discharge_based_run(kernel)
except NotImplementedError:
    run_on_gpu(kernel)  # compiled path

When it happens

Trigger: Running a kernel containing inline_mgpu/inline_ptx through discharge-based paths: interpretation on CPU, pallas discharge mode, or APIs that require discharging the kernel.

Common situations: Trying to debug a Pallas kernel with the interpreter; using a pipeline that discharges transforms (e.g. certain checkpointing/remat paths).

Related errors

• DMA partial discharge add=True not yet implemented.
• Explicit sharding is not currently supported in the pallas-t
• group_offset is not currently supported in the pallas-triton
• get not supported yet
• get not supported yet for block shape {b}

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