{"record":{"id":"4773a74de3e40cf4","repo":"jax-ml/jax","slug":"megacore-mode-must-be-one-of-kv-head-batch","errorCode":null,"errorMessage":"megacore_mode must be one of ['kv_head', 'batch', None]","messagePattern":"megacore_mode must be one of \\['kv_head', 'batch', None\\]","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/paged_attention/paged_attention_kernel.py","lineNumber":483,"sourceCode":"    raise ValueError(\n        f\"The dtype of `lengths` must be int32. Got {lengths.dtype}\"\n    )\n\n  # TODO(dinghua): get the actual cores per chip once there's an official API.\n  if megacore_mode == \"kv_head\":\n    if num_kv_heads % 2 != 0:\n      raise ValueError(\n          \"number of KV heads must be even when megacore_mode is 'kv_head'\"\n      )\n    num_cores = 2\n  elif megacore_mode == \"batch\":\n    if batch_size % 2 != 0:\n      raise ValueError(\"batch size must be even when megacore_mode is 'batch'\")\n    num_cores = 2\n  elif megacore_mode is None:\n    num_cores = 1\n  else:\n    raise ValueError(\"megacore_mode must be one of ['kv_head', 'batch', None]\")\n\n  num_groups = num_q_heads // num_kv_heads\n  if (num_groups) % 8 != 0:\n    # Reshape q to hint XLA to pick a <1x128> layout otherwise it will pick a\n    # <8x128> layout for a <1x128> memref inside the kernel and error out.\n    q = q.reshape(batch_size, num_q_heads, 1, head_dim)\n    if megacore_mode == \"kv_head\":\n      q_block_spec = pl.BlockSpec(\n          (None, num_groups, None, head_dim),\n          lambda core_index, b, h, *_: (b, h * num_cores + core_index, 0, 0),\n      )\n    elif megacore_mode == \"batch\":\n      q_block_spec = pl.BlockSpec(\n          (None, num_groups, None, head_dim),\n          lambda core_index, b, h, *_: (b * num_cores + core_index, h, 0, 0),\n      )\n    else:\n      q_block_spec = pl.BlockSpec(","sourceCodeStart":465,"sourceCodeEnd":501,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/paged_attention/paged_attention_kernel.py#L465-L501","documentation":"megacore_mode selects how work is partitioned across the TPU Matmul cores and only accepts the literal strings 'kv_head', 'batch', or None. Any other value (including case variants or typo'd strings) falls through to the final else branch and raises this error.","triggerScenarios":"Passing megacore_mode='KV_HEAD', 'batch_head', 'none' (string), 'kv', or an unset config variable defaulting to '' instead of None.","commonSituations":"Config-driven megacore selection from YAML/env vars where None is serialized as the string 'none'; typos in copied sample code; API changes from older versions that accepted other modes.","solutions":["Use exactly one of 'kv_head', 'batch', or None (Python None, not 'none')","Normalize config input: mode = None if mode in ('', 'none', 'None') else mode","Add a unit test asserting megacore_mode membership in the allowed set"],"exampleFix":"// before\npaged_attention(..., megacore_mode='none')\n// after\npaged_attention(..., megacore_mode=None)","handlingStrategy":"type-guard","validationCode":"ALLOWED = ('kv_head', 'batch', None)\nmegacore_mode = None if megacore_mode in ('', 'none', 'None') else megacore_mode\nassert megacore_mode in ALLOWED","typeGuard":"def is_valid_megacore_mode(m): return m in ('kv_head', 'batch', None)","tryCatchPattern":null,"preventionTips":["Never source megacore_mode from raw config strings without normalization","Use Python None, not the string 'none'"],"tags":["jax","pallas","tpu","paged-attention","megacore","invalid-argument"],"backgroundTag":"invalid-argument-value","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}