{"record":{"id":"cda5f1d2418fbfa7","repo":"jax-ml/jax","slug":"4-bit-block-scaled-mma-only-supports-k-fastest-ope","errorCode":null,"errorMessage":"4-bit block scaled MMA only supports K-fastest operands, but A is M-fastest","messagePattern":"4-bit block scaled MMA only supports K-fastest operands, but A is M-fastest","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":546,"sourceCode":"    a_fastest = mma_utils.Dim.K\n    a_k_instr_strides = None\n    a_m_group_stride = a_k_group_stride = a_desc_base = None\n  (\n      (b_desc_base, b_k_instr_strides),\n      (b_n_group_stride, b_k_group_stride),\n      b_fastest,\n  ) = mma_utils.create_descriptor(\n      b,\n      swizzle=b_swizzle,\n      group_size=(k_group_elems, n_group_elems),\n      logical_k_major=True,\n      mma_bytewidth_k=64 if is_sparse else 32,\n      split_const=True,\n  )\n\n  if is_scaled and utils.bitwidth(mma_a_element_type) == 4:\n    if a_fastest != mma_utils.Dim.K:\n      raise ValueError(\n          \"4-bit block scaled MMA only supports K-fastest operands, but A is M-fastest\"\n      )\n    if b_fastest != mma_utils.Dim.K:\n      raise ValueError(\n          \"4-bit block scaled MMA only supports K-fastest operands, but B is N-fastest\"\n      )\n  if is_sparse:\n    if b_swizzle == 32 and b_fastest == mma_utils.Dim.K:\n      raise NotImplementedError(\n          \"B tiling too small. Increase swizzle or transpose the input.\"\n      )\n\n  # Step 4. Issue the instructions.\n  true = arith.constant(ir.IntegerType.get_signless(1), 1)\n  n_collective_group_elems = n_group_elems * num_cta\n  n_col_groups = n_groups // n_lane_groups\n  assert d.layout.base_tile_shape[0] % 4 == 0\n  lanes_per_n_group = d.layout.base_tile_shape[0] // 4","sourceCodeStart":528,"sourceCodeEnd":564,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L528-L564","documentation":"4-bit block-scaled MMA (e.g. MXFP4/NVFP4) hardware only supports operands whose fastest-varying dimension is K. If the A operand is laid out M-fastest, mma raises this error.","triggerScenarios":"Calling mma() with is_scaled=True, 4-bit A, and a layout where a_fastest is Dim.M (e.g. a transposed or col-major A).","commonSituations":"Reusing 8-bit kernel layouts for MXFP4; transposing A for coalescing purposes which flips the fastest dim.","solutions":["Transpose A so K is the fastest-varying dimension and update swizzle accordingly","Use an 8-bit scaled type if an M-fastest A layout is required"],"exampleFix":"# before\na = TensorMemRefView(buf, (m, k), dt, layout=col_major)  # M-fastest\ntcgen05.mma(a, b, d, a_scale=asc, b_scale=bsc, scale_block=16)\n# after\na = TensorMemRefView(buf, (m, k), dt, layout=row_major)  # K-fastest\ntcgen05.mma(a, b, d, a_scale=asc, b_scale=bsc, scale_block=16)","handlingStrategy":"validation","validationCode":"assert a_fastest == mma_utils.Dim.K, '4-bit scaled MMA needs K-fastest A'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Use K-major (row-major) A layouts for all 4-bit scaled kernels","Add layout assertions in kernel prologue"],"tags":["gpu","mosaic","tcgen05","layout","mxfp4","block-scaling"],"backgroundTag":"invalid-memory-layout","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}