{"record":{"id":"0bfc98b9663e8f07","repo":"jax-ml/jax","slug":"only-m-128-and-m-64-are-supported-for-mma-but-got","errorCode":null,"errorMessage":"Only M=128 and M=64 are supported for MMA, but got M={m}","messagePattern":"Only M=128 and M=64 are supported for MMA, but got M=(.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":281,"sourceCode":"    # 2CTA M=128 instruction uses a different TMEM layout than 1CTA M=64.\n    expected_d_layout = _infer_tmem_layout(d.shape, collective, packing=1)\n    if d.layout != expected_d_layout:\n      raise ValueError(\n          f\"Accumulator layout mismatch: expected {expected_d_layout}, got {d.layout}\"\n      )\n    if collective:\n      n_lane_groups = 1\n    else:\n      n_lane_groups = 2\n      # We can't split N into groups if we would partition it below the tile size.\n      # TODO: We only need to check this if N is the minormost dim in B.\n      if 8 * b_swizzle // utils.bitwidth(a_element_type) > n // n_lane_groups:\n        raise ValueError(\n            f\"Swizzle={b_swizzle} is too big for MMA with M=64. Try\"\n            \" lowering it.\"\n        )\n  else:\n    raise ValueError(f\"Only M=128 and M=64 are supported for MMA, but got M={m}\")\n  f32 = ir.F32Type.get()\n  f16 = ir.F16Type.get()\n  s32 = ir.IntegerType.get_signless(32)\n  elem_type_str = (\n      f\"{a_element_type}\"\n      if a_element_type == b_element_type\n      else f\"({a_element_type}, {b_element_type})\"\n  )\n  if a_element_type == f32 or a_element_type == ir.BF16Type.get():\n    if a_element_type == f32 and is_sparse:\n      raise NotImplementedError(\"Sparse MMA unsupported for f32\")\n    if is_scaled:\n      raise ValueError(\n          f\"MMA with element type {elem_type_str} does not support block scaling\"\n      )\n    if d.dtype != f32:\n      raise ValueError(\n          f\"MMA with element type {elem_type_str} only supports accumulators\"","sourceCodeStart":263,"sourceCodeEnd":299,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L263-L299","documentation":"The Blackwell tcgen05 MMA instruction only supports M=128 and M=64 operand shapes. Any other leading dimension of the A operand is rejected up front with this ValueError.","triggerScenarios":"Calling tcgen05.mma where lhs's first dimension is anything other than 128 or 64 (e.g. 256, 32, or a symbolic/dynamic M).","commonSituations":"Using tile sizes inherited from Hopper wgmma kernels (e.g. M=64 works but M=32 or M=256 does not); computing M from a dynamic batch dimension that resolves to an unsupported value; off-by-one/slicing bugs that shrink M.","solutions":["Set the A operand's leading dimension to 128 or 64 (pad the tile if the natural M differs)","If M is dynamic, assert or clamp it to 128/64 before building the MMA","Check for slicing mistakes that produced an unexpected M"],"exampleFix":"# before\nm = lhs.shape[0]  # e.g. 32\ntcgen05.mma(lhs, rhs, d)\n# after\nassert lhs.shape[0] in (64, 128), f\"unsupported M={lhs.shape[0]}\"\ntcgen05.mma(lhs, rhs, d)","handlingStrategy":"validation","validationCode":"assert lhs.shape[0] in (64, 128), f'M must be 64 or 128, got {lhs.shape[0]}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Clamp or pad M to 64/128 before building the MMA","Fail fast on dynamic M values at trace time"],"tags":["jax","mosaic","gpu","mma","shape-validation","tcgen05"],"backgroundTag":"unsupported-matrix-dimension","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}