{"record":{"id":"76ad9e47aa02b726","repo":"jax-ml/jax","slug":"4-bit-block-scaled-mma-only-supports-k-fastest-ope-76ad9e","errorCode":null,"errorMessage":"4-bit block scaled MMA only supports K-fastest operands, but B is N-fastest","messagePattern":"4-bit block scaled MMA only supports K-fastest operands, but B is N-fastest","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":550,"sourceCode":"      (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\n  a_sparse_addr_base = a_sparse_metadata.address if is_sparse else None\n  a_scale_addr_base = a_scale.address if is_scaled else None  # pyrefly: ignore[missing-attribute]\n  b_scale_addr_base = b_scale.address if is_scaled else None  # pyrefly: ignore[missing-attribute]\n  # B scales are padded when N is short, so it can't be derived from n_collective_group_elems.","sourceCodeStart":532,"sourceCodeEnd":568,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L532-L568","documentation":"4-bit block-scaled MMA requires the B operand to be K-fastest as well. If B's fastest-varying dimension is N, mma raises this error.","triggerScenarios":"Calling mma() with is_scaled=True, 4-bit operands, and b_fastest == Dim.N (B stored row-major over (N, K) transposed).","commonSituations":"Standard GEMM weight layout (N-major B) reused from an 8-bit kernel; missing transpose of the weights for NVFP4.","solutions":["Store/transpose B so K is the fastest dimension","Adjust b_swizzle to a value compatible with the K-fastest layout"],"exampleFix":"# before\nb = weights  # N-fastest\ntcgen05.mma(a, b, d, a_scale=asc, b_scale=bsc, scale_block=16)\n# after\nb_t = transpose_to_k_fastest(weights)\ntcgen05.mma(a, b_t, d, a_scale=asc, b_scale=bsc, scale_block=16)","handlingStrategy":"validation","validationCode":"assert b_fastest == mma_utils.Dim.K, '4-bit scaled MMA needs K-fastest B'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Pre-transpose weights to K-major at quantization time","Assert b_fastest for 4-bit scaled paths"],"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"}