{"record":{"id":"fbf99367adea7c51","repo":"jax-ml/jax","slug":"sparse-mma-not-supported-for-m-64","errorCode":null,"errorMessage":"Sparse MMA not supported for M=64","messagePattern":"Sparse MMA not supported for M=64","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":261,"sourceCode":"        \"MMA requires A and B to have the same element type, except that \"\n        \"FP8 types (f8E4M3FN and f8E5M2) may be mixed; got: \"\n        f\"{a_element_type} and {b_element_type}\"\n    )\n  if d.shape != (m, n * num_cta):\n    raise ValueError(\n        f\"Accumulator shape mismatch: expected {(m, n * num_cta)}, got {d.shape}\"\n    )\n  if m == 128:\n    if d.layout != (expected_d_layout := tmem_default_layout(packing=1)):\n      raise ValueError(\n          f\"Accumulator layout mismatch: expected {expected_d_layout}, got {d.layout}\"\n      )\n    n_lane_groups = 1\n  elif m == 64:\n    if is_scaled and not collective:\n      raise NotImplementedError(\"MMA with block scaling is not supported for 1CTA M=64\")\n    if is_sparse:\n      raise NotImplementedError(\"Sparse MMA not supported for M=64\")\n    # Watch out: this layout must be consistent with A's layout (up to packing).\n    # 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        )","sourceCodeStart":243,"sourceCodeEnd":279,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L243-L279","documentation":"Raised by Mosaic GPU's tcgen05 mma op when a sparse MMA is requested with M=64. The underlying tcgen05 sparse tensor-core instruction only exists for the M=128 shape, so the library explicitly rejects the combination rather than emitting invalid MLIR.","triggerScenarios":"Calling jax.experimental.mosaic.gpu.tcgen05.mma with a lhs whose leading (M) dimension is 64 while passing a sparse metadata / is_sparse path (e.g. building a sparse warp-level MMA pipeline on Blackwell).","commonSituations":"Porting a sparse kernel originally written for M=128 tiles to smaller 64-wide tiles; reusing a sparse GEMM template and changing only the tile size; targeting sm_100/sm_103 with truncated M.","solutions":["Change the M dimension of the A operand/accumulator to 128 (pad the tile if necessary)","If M must stay 64, drop the sparse path and use a dense MMA instead","Split your workload so sparse MMAs are issued only on 128-row tiles and dense MMAs handle the remainder"],"exampleFix":"# before\nacc = tcgen05.mma(lhs_64, rhs, acc, sparse_metadata=meta)  # M=64 -> raises\n# after\nlhs_128 = pad_tile_to_m128(lhs_64)\nacc = tcgen05.mma(lhs_128, rhs, acc, sparse_metadata=meta)","handlingStrategy":"validation","validationCode":"m = lhs.shape[0]\nassert not (is_sparse and m == 64), 'Sparse tcgen05 MMA requires M=128'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep sparse paths parameterized by tile size and assert M==128 at kernel entry","Unit-test kernels at every supported tile shape"],"tags":["jax","mosaic","gpu","sparse","mma","tensor-core","tcgen05","blackwell"],"backgroundTag":"unsupported-operation-shape","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}