{"record":{"id":"18cc2d4754251f00","repo":"jax-ml/jax","slug":"can-only-perform-collective-mma-along-minormost-cl","errorCode":null,"errorMessage":"Can only perform collective MMA along minormost cluster axis. Got {collective_axis}, expected {minormost_cluster_axis}.","messagePattern":"Can only perform collective MMA along minormost cluster axis\\. Got (.+?), expected (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":3195,"sourceCode":"  with predicate_ctx:\n    mgpu.dialect.tcgen05_commit_arrive(\n        barrier_ref.as_barrier_memref(), collective=collective\n    )\n  return []\n\n\ndef _collective_mma_predicate(ctx: lowering.LoweringRuleContext,\n                              collective_axis: str) -> ir.Value:\n  \"\"\"Computes a predicate to run only on the leader block.\"\"\"\n  cluster_axis = lowering._resolve_cluster_axis(\n      ctx.module_ctx.axis_names, collective_axis)\n  if cluster_axis != gpu_dialect.Dimension(0):\n    # Note: resolve_cluster_axis checks if axis_names exists.\n    assert ctx.module_ctx.axis_names is not None\n    if len(ctx.module_ctx.axis_names.cluster) <= 1:\n      raise ValueError(\"No cluster axes found.\")\n    minormost_cluster_axis = ctx.module_ctx.axis_names.cluster[0]\n    raise ValueError(\n        \"Can only perform collective MMA along minormost cluster axis. \"\n        f\"Got {collective_axis}, expected {minormost_cluster_axis}.\")\n  index = ir.IndexType.get()\n  is_leader_block = arith_dialect.cmpi(\n      arith_dialect.CmpIPredicate.eq,\n      mgpu.utils.cluster_idx(cluster_axis), mgpu.c(0, index))\n  return is_leader_block\n\n\ncommit_tmem_p = jax_core.Primitive(\"commit_tmem\")\ncommit_tmem_p.multiple_results = True\n\n\n@commit_tmem_p.def_effectful_abstract_eval\ndef _commit_tmem_abstract_eval():\n  return (), {gpu_core._memory_effect}\n\n","sourceCodeStart":3177,"sourceCodeEnd":3213,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L3177-L3213","documentation":"Collective MMA can only multicast along the minormost (first) cluster axis because that is the axis the hardware multicast uses. If the requested collective_axis resolves to a non-zero cluster dimension, the lowering rejects it and tells you which axis it expected.","triggerScenarios":"Passing collective_axis naming a cluster axis that is not axis_names.cluster[0], e.g. declaring cluster axes ('x','y') and requesting collective_axis='y'.","commonSituations":"Reordering BlockSpec/collective axes; assuming any cluster axis can be multicast along.","solutions":["Reorder your axes so the requested collective axis is the minormost cluster axis (first in cluster list)","Change collective_axis to name the first cluster axis"],"exampleFix":"// before\n# cluster axes declared as ('x','y'); requesting 'y'\nspec = BlockSpec(..., collective_axis='y')\n// after\nspec = BlockSpec(..., collective_axis='x')","handlingStrategy":"validation","validationCode":"assert collective_axis == cluster_axis_names[0]","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Put the multicast axis first in the cluster axis list","Document axis ordering in kernel launch config"],"tags":["jax","pallas","tcgen05","cluster","collective-mma","axis-order"],"backgroundTag":"axis-order-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}