jax-ml/jax · error · ValueError

Need to specify backward blocks.

Error message

Need to specify backward blocks.

What it means

The splash attention backward primitive requires explicit backward block sizes; unlike some kernels they are not inferred. Calling jax.grad (or .with_backward()) on a splash kernel whose BlockSizes lacks backward blocks raises this immediately.

Source

Thrown at jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py:2279

    residual_checkpoint_name: str | None,
    mask_function: MaskFunctionType | None,
    attn_logits_soft_cap: float | None,
    interpret: bool,
    res: SplashResidualsType,
    do: jax.Array,
) -> tuple[
    mask_info_lib.MaskInfo | None,  # fwd_mask_info
    mask_info_lib.MaskInfo | None,  # dq_mask_info
    mask_info_lib.MaskInfo | None,  # dvk_mask_info
    jax.Array,  # q
    jax.Array,  # k
    jax.Array,  # v
    SegmentIds | None,  # segmend_ids
    jax.Array | None,  # sinks
]:
  del save_residuals, residual_checkpoint_name
  if not block_sizes.has_backward_blocks:
    raise ValueError("Need to specify backward blocks.")
  bq_dq, bkv_dq = block_sizes.block_q_dq, block_sizes.block_kv_dq
  bq_dkv, bkv_dkv_memory, bkv_dkv_compute = (
      block_sizes.block_q_dkv,
      block_sizes.block_kv_dkv,
      block_sizes.block_kv_dkv_compute,
  )
  use_fused_bwd_kernel = block_sizes.use_fused_bwd_kernel
  (
      q,
      k,
      v,
      segment_ids,
      sinks,
      o,
      logsumexp,
      dq_mask_info,
      dkv_mask_info,
  ) = res

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Create the kernel via BlockSizes.fwd_bwd(...) or call .with_backward(...) on the forward kernel so backward blocks are populated
2. Pass explicit block_q_dq/block_kv_dq/block_q_dkv/block_kv_dkv when constructing BlockSizes for training
3. If you only need inference, avoid jax.grad over this function

Example fix

// before
kernel = make_splash_attention(mask, block_sizes=BlockSizes(...fwd only...))
loss = jax.grad(fn)(...)  # raises
// after
kernel = make_splash_attention(mask, block_sizes=BlockSizes.fwd_bwd(...))
loss = jax.grad(fn)(...)

assert block_sizes.has_backward_blocks, 'need BlockSizes.fwd_bwd(...) for grad'

When it happens

Trigger: Calling make_splash_attention(...) without with_backward(...) (or without passing backward BlockSizes) and then differentiating the function, which routes into _splash_attention_bwd.

Common situations: Reusing a forward-only kernel config for training; upgrading jax versions where defaults for backward blocks changed from implicit to explicit; constructing BlockSizes manually and leaving bwd fields None.

Related errors

• {bkv=} must be a multiple of {NUM_LANES}.
• {bkv_compute=} should not be greater than {bkv=}
• {bkv=} should be a multiple of {bkv_compute=}
• Higher-order AD not supported.
• Block sizes for dq kernel are not needed with a fused kernel

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