jax-ml/jax · error · ValueError

Kernel output {j} in HBM but does not have trivial BlockSpec

Error message

Kernel output {j} in HBM but does not have trivial BlockSpec.

What it means

Symmetric to the input case: in TPU interpret mode, an output residing in HBM must have a trivial BlockSpec whose block shape equals the full output shape. If the kernel's block shape for that output differs from the output tensor shape, this ValueError is raised naming the output index.

Source

Thrown at jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py:2432

              functools.partial(store, output_name=output_names[index]),
              TOKEN_SHAPE_DTYPE,
              token,
              device_id,
              core_index,
              TPU_MEMORY_SPACE_IDXS[mosaic_core.MemorySpace.HBM],
              output_buffer_ids[index],
              (transform,),
              kernel_output_val,
              cur_block_indices[num_inputs + index],
              grid_point,
          )
          return token

        output_slices : list[Any] = []
        for j, var in enumerate(output_vars):
          if output_var_memory_spaces[j] is _HBM:
            if var.aval.shape != block_shapes[num_inputs + j]:
              raise ValueError(
                  f'Kernel output {j} in HBM but does not have trivial'
                  ' BlockSpec.'
              )
            output_slices.append(None)
            continue
          assert len(cur_start_indices[num_inputs + j].shape) == 1
          assert len(next_start_indices[num_inputs + j].shape) == 1
          transform = indexing.NDIndexer(
              indices=tuple(
                  indexing.ds(st, sz) if not iid else st
                  for st, sz, iid in zip(
                      cur_start_indices[num_inputs + j],
                      block_shapes[num_inputs + j],
                      is_squeeze_dim[num_inputs + j],
                  )
              ),
              shape=output_vals[j].shape,
              int_indexer_shape=(),

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use a trivial BlockSpec (full shape, zero start indices) for HBM outputs
2. Write results via async_copy from a VMEM buffer to the HBM output instead of blocked stores
3. Split outputs so blocked ones are VMEM-intermediates copied back manually

Example fix

# before
out_spec = BlockSpec((BM, BN), lambda i, j: (i*BM, j*BN))  # HBM output blocked
# after
out_spec = BlockSpec((M, N), lambda i, j: (0, 0))  # trivial for HBM output

for j, (o, spec) in enumerate(zip(out_shapes, out_specs)):
    if output_memory_spaces[j] == 'hbm' and spec is not None:
        assert tuple(spec.block_shape) == tuple(o.shape), f'output {j} in HBM needs trivial BlockSpec'

def has_trivial_blockspec(spec, shape) -> bool:
    return spec is None or tuple(spec.block_shape) == tuple(shape)

try:
    interpret_run(kernel)
except ValueError as e:
    if 'does not have trivial BlockSpec' in str(e) and 'output' in str(e):
        # use full-shape BlockSpec for that output and retry
        raise

When it happens

Trigger: Declaring an HBM output with a blocked BlockSpec (block shape smaller than out shape) in a kernel run under TPU interpret mode.

Common situations: Streaming-style kernels writing blocks directly to HBM outputs; reusing one grid/BlockSpec tuple for inputs and outputs where only inputs should be blocked; GPU Pallas kernels ported to TPU interpretation.

Related errors

• Kernel input {j} in HBM but does not have trivial BlockSpec.
• Out-of-bounds read of ({device_id} {local_core_id} {memory_s
• Out-of-bounds block index {block_indices} for input "{input_
• Out-of-bounds block index {block_indices} for output "{outpu
• {primitive_name}: Buffers with a memory space of HBM or ANY

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