jax-ml/jax · error · ValueError

Block sizes for dq kernel are not needed with a fused kernel

Error message

Block sizes for dq kernel are not needed with a fused kernel.

What it means

Splash attention's BlocksConfig supports a fused backward kernel (use_fused_bwd_kernel=True); in that mode the dq block sizes (block_q_dq, block_kv_dq) are determined internally, so specifying them is contradictory and __post_init__ raises ValueError.

Source

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

  block_kv_dkv_compute: int | None = None

  block_q_dq: int | None = None
  block_kv_dq: int | None = None

  use_fused_bwd_kernel: bool = False

  q_layout: QKVLayout = QKVLayout.HEAD_DIM_MINOR
  k_layout: QKVLayout = QKVLayout.HEAD_DIM_MINOR
  v_layout: QKVLayout = QKVLayout.HEAD_DIM_MINOR

  def __post_init__(self):
    if self.block_kv_compute is None:
      object.__setattr__(self, "block_kv_compute", self.block_kv)
    if self.block_kv_dkv_compute is None:
      object.__setattr__(self, "block_kv_dkv_compute", self.block_kv_dkv)
    if self.use_fused_bwd_kernel:
      if self.block_q_dq is not None or self.block_kv_dq is not None:
        raise ValueError(
            "Block sizes for dq kernel are not needed with a fused kernel."
        )

  @property
  def has_backward_blocks(self) -> bool:
    backward_blocks = (
        self.block_q_dkv, self.block_kv_dkv, self.block_kv_dkv_compute,
    )
    if not self.use_fused_bwd_kernel:
      backward_blocks += (self.block_q_dq, self.block_kv_dq)
    return all(b is not None for b in backward_blocks)

  @classmethod
  def get_default(cls):
    # TODO(apaszke,sharadmv): Select better parameters based on a heuristic.
    return BlockSizes(
        block_q=128,
        block_kv=128,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set block_q_dq=None and block_kv_dq=None when enabling the fused kernel
2. Or keep the dq block sizes and leave use_fused_bwd_kernel=False
3. Sanitize configs: if fused, drop dq fields before constructing the dataclass

Example fix

// before
BlocksConfig(..., use_fused_bwd_kernel=True, block_q_dq=128, block_kv_dq=64)
// after
BlocksConfig(..., use_fused_bwd_kernel=True, block_q_dq=None, block_kv_dq=None)

if cfg.get('use_fused_bwd_kernel'):
    cfg['block_q_dq'] = None
    cfg['block_kv_dq'] = None
blocks = BlocksConfig(**cfg)

def is_valid_blocks_config(cfg) -> bool:
    return not (cfg.get('use_fused_bwd_kernel') and (cfg.get('block_q_dq') is not None or cfg.get('block_kv_dq') is not None))

When it happens

Trigger: Constructing a BlocksConfig with use_fused_bwd_kernel=True while also setting block_q_dq or block_kv_dq to non-None values.

Common situations: Copying a config from a non-fused setup (where dq block sizes were tuned) and then flipping use_fused_bwd_kernel=True without clearing the dq fields; merging configs programmatically.

Related errors

• Number of cores or threads must be at least 1, but got {self
• Vector clock size must be at least 1, but got {self.vector_c
• {sliding_window=} must be positive.
• {soft_cap=} must not be 0.0.
• Higher-order AD not supported.

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