{"record":{"id":"da70b097539c1418","repo":"jax-ml/jax","slug":"block-sizes-for-dq-kernel-are-not-needed-with-a-fu","errorCode":null,"errorMessage":"Block sizes for dq kernel are not needed with a fused kernel.","messagePattern":"Block sizes for dq kernel are not needed with a fused kernel\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py","lineNumber":523,"sourceCode":"  block_kv_dkv_compute: int | None = None\n\n  block_q_dq: int | None = None\n  block_kv_dq: int | None = None\n\n  use_fused_bwd_kernel: bool = False\n\n  q_layout: QKVLayout = QKVLayout.HEAD_DIM_MINOR\n  k_layout: QKVLayout = QKVLayout.HEAD_DIM_MINOR\n  v_layout: QKVLayout = QKVLayout.HEAD_DIM_MINOR\n\n  def __post_init__(self):\n    if self.block_kv_compute is None:\n      object.__setattr__(self, \"block_kv_compute\", self.block_kv)\n    if self.block_kv_dkv_compute is None:\n      object.__setattr__(self, \"block_kv_dkv_compute\", self.block_kv_dkv)\n    if self.use_fused_bwd_kernel:\n      if self.block_q_dq is not None or self.block_kv_dq is not None:\n        raise ValueError(\n            \"Block sizes for dq kernel are not needed with a fused kernel.\"\n        )\n\n  @property\n  def has_backward_blocks(self) -> bool:\n    backward_blocks = (\n        self.block_q_dkv, self.block_kv_dkv, self.block_kv_dkv_compute,\n    )\n    if not self.use_fused_bwd_kernel:\n      backward_blocks += (self.block_q_dq, self.block_kv_dq)\n    return all(b is not None for b in backward_blocks)\n\n  @classmethod\n  def get_default(cls):\n    # TODO(apaszke,sharadmv): Select better parameters based on a heuristic.\n    return BlockSizes(\n        block_q=128,\n        block_kv=128,","sourceCodeStart":505,"sourceCodeEnd":541,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py#L505-L541","documentation":"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.","triggerScenarios":"Constructing a BlocksConfig with use_fused_bwd_kernel=True while also setting block_q_dq or block_kv_dq to non-None values.","commonSituations":"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.","solutions":["Set block_q_dq=None and block_kv_dq=None when enabling the fused kernel","Or keep the dq block sizes and leave use_fused_bwd_kernel=False","Sanitize configs: if fused, drop dq fields before constructing the dataclass"],"exampleFix":"// before\nBlocksConfig(..., use_fused_bwd_kernel=True, block_q_dq=128, block_kv_dq=64)\n// after\nBlocksConfig(..., use_fused_bwd_kernel=True, block_q_dq=None, block_kv_dq=None)","handlingStrategy":"validation","validationCode":"if cfg.get('use_fused_bwd_kernel'):\n    cfg['block_q_dq'] = None\n    cfg['block_kv_dq'] = None\nblocks = BlocksConfig(**cfg)","typeGuard":"def is_valid_blocks_config(cfg) -> bool:\n    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))","tryCatchPattern":null,"preventionTips":["When flipping use_fused_bwd_kernel=True, clear block_q_dq/block_kv_dq in the same change","Centralize BlocksConfig construction behind a factory that enforces mutual exclusivity"],"tags":["jax","pallas","tpu","splash-attention","config-validation"],"backgroundTag":"conflicting-config-options","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}