jax-ml/jax · error · ValueError

Only workgroup memory is supported but got {ref}.

Error message

Only workgroup memory is supported but got {ref}.

What it means

During Pallas GPU lowering, a dtype bitcast was requested on a ref that does not live in workgroup (shared) memory. _handle_dtype_bitcast checks mgpu_utils.is_smem_ref and only knows how to reinterpret SMEM memrefs; TMEM, WMEM, or global-memory refs cannot be bitcast here.

Source

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

    ref: the reference to bitcast.
    src_dtype: the source element type.
    dst_dtype: the destination element type.

  Returns:
    A bitcasted version of `ref` with element type `dst_dtype`.

  Raises:
    ValueError: if the source ref is not in SMEM.
  """
  if src_dtype == dst_dtype:
    return ref
  if src_dtype != ir.IntegerType.get_signless(8):
    raise NotImplementedError(
        "Data type bitcast is only supported from i8 to other types."
    )
  ref_ty = ir.MemRefType(ref.type)
  if not mgpu_utils.is_smem_ref(ref_ty):
    raise ValueError(f"Only workgroup memory is supported but got {ref}.")
  if len(ref_ty.shape) != 1:
    raise NotImplementedError(
        "Data type bitcast is only supported for 1D arrays."
    )
  [stride], _ = ref_ty.get_strides_and_offset()
  if stride != 1:
    raise ValueError(
        "Data type bitcast is only supported for contiguous 1D arrays, but got "
        f"stride={stride}."
    )
  [shape_bytes] = ref_ty.shape
  shape_bitwidth = shape_bytes * 8
  target_bitwidth = mgpu_utils.bitwidth(dst_dtype)

  if shape_bitwidth % target_bitwidth:
    raise ValueError(
        f"Can not bitcast memory region of size {shape_bitwidth} bits to dtype "
        f"with {target_bitwidth} bits."

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Move the aliased buffer to SMEM (use regular SMEM scratch instead of TMEM/WMEM)
2. Copy the data from the non-SMEM ref into an SMEM i8 buffer, then bitcast that buffer
3. Convert values explicitly with jax.lax.bitcast_convert after loading

Example fix

// before
buf = alloc_tmem(...)  # or non-SMEM ref
view = buf.view(jnp.float32)
// after
buf = alloc_smem(...)  # SMEM-backed i8 scratch
view = buf.view(jnp.float32)

# only alias refs allocated in SMEM
assert buf.memory_space == 'smem', 'bitcast alias requires an SMEM buffer'

def is_smem_backed(buf) -> bool:
    return getattr(buf, 'memory_space', 'smem') == 'smem'

When it happens

Trigger: Aliasing a Ref stored in tensor memory (tcgen05.TMEMRef) or non-SMEM memory with a different dtype, causing _extract_aliased_ref to route through the bitcast path on a non-SMEM ref.

Common situations: Mixing tmem (tcgen05) buffers with dtype reinterpretation on Hopper/Blackwell; kernels migrated from SMEM scratch to TMEM allocations while keeping view/alias dtype tricks.

Related errors

• Only byte-aligned bitcasts are supported.
• packed, collective and layout arguments are only supported f
• Only SMEM and TMEM refs are supported.
• All aliased Refs must have the same memory space (SMEM or TM
• Data type bitcast is only supported from i8 to other types.

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