jax-ml/jax · error · NotImplementedError

Buffer poisoning is not supported on GPU yet.

Error message

Buffer poisoning is not supported on GPU yet.

What it means

Buffer poisoning (a debug mode that fills output buffers with NaNs/garbage to catch uses of uninitialized memory, enabled via JAX flag pallas_default_buffer_poisoning / analogous config) is not implemented for the GPU Pallas backend, so run_scoped lowering refuses to proceed when it is on.

Source

Thrown at jax/_src/pallas/mosaic_gpu/lowering.py:3567

    raise NotImplementedError(
        "debug_print only supports printing of scalar values, or a single array"
        " value when using the Mosaic GPU backend."
    )

  return ()


@register_lowering_rule(primitives.run_scoped_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(primitives.run_scoped_p, mgpu.LoweringSemantics.Warpgroup)
def _run_scoped_lowering_rule(
    ctx: LoweringRuleContext,
    *consts,
    jaxpr: jax_core.Jaxpr,
    collective_axes,
    **_,
):
  if pallas_core.poison_buffers_enabled():
    raise NotImplementedError("Buffer poisoning is not supported on GPU yet.")
  input_refs = []
  should_discharge = []
  wg_axis = ctx.module_ctx.axis_names.wg
  is_multithreaded = wg_axis is not None
  is_thread_collective = is_multithreaded and collective_axes == (wg_axis,)
  # Make sure everyone has exited previous scoped allocations. Note that we
  # don't synchronize when we exit the allocation, but only when we might want
  # to reuse its memory again.
  if collective_axes and collective_axes != (wg_axis,):
    raise ValueError(
        "Only thread-collective allocations are supported in run_scoped."
    )
  if is_multithreaded and is_thread_collective:
    gpu_dialect.barrier()
  with contextlib.ExitStack() as alloc_stack:
    for v in jaxpr.invars:
      aval = cast(ShapedAbstractValue, v.aval)
      if isinstance(aval, gpu_core.WGMMAAbstractAccumulatorRef):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Disable buffer poisoning when running GPU kernels (set the poisoning value to None / off)
2. Gate the flag by platform: only enable for TUSA/TPU paths
3. Use other uninitialized-memory checks (e.g. compute-sanitizer) on GPU instead

Example fix

// before
jax.config.update('jax_default_buffer_poisoning_value', float('nan'))
// after  # GPU
jax.config.update('jax_default_buffer_poisoning_value', None)

import jax
if jax.config.jax_default_buffer_poisoning_value is not None and device == 'gpu':
    jax.config.update('jax_default_buffer_poisoning_value', None)

When it happens

Trigger: Enabling buffer poisoning (e.g. jax.config.update('jax_default_buffer_poisoning_value', ...) or the pallas poisoning env/flag) and running a GPU kernel that uses pallas.run_scoped.

Common situations: Turning on aggressive memory debugging flags globally and then running the same suite on GPU; CI configs that enable poisoning on TPU-only.

Related errors

• This gmm kernel only supports either (m, k) x (g, k, n) -> (
• Group sizes {group_sizes.shape=} must match first dimension
• group_offset is not currently supported in the pallas-triton
• Compiler params for platform {platform} cannot be used for {
• Memory space {self.memory_space} is not supported by mesh {s

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