{"record":{"id":"ec5e5b033c05294e","repo":"jax-ml/jax","slug":"b-scale-shape-0-must-be-a-multiple-of-128-and","errorCode":null,"errorMessage":"B scale shape[0] must be a multiple of 128 and >= N={n * num_cta}, got {b_scale.shape[0]}","messagePattern":"B scale shape\\[0\\] must be a multiple of 128 and >= N=(.+?), got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":487,"sourceCode":"          f\" {a_scale.shape}\"\n      )\n    if a_scale.layout != scales_layout():\n      raise ValueError(f\"A scale layout {a_scale.layout} is not supported\")\n    if collective and m == 64:\n      if b_scale.layout != b_scales_m64_collective_layout():\n        raise ValueError(\n            \"Expected B scales to have a M=64 collective layout, got\"\n            f\" {b_scale.layout}\"\n        )\n    elif m == 128:\n      if b_scale.layout != scales_layout():\n        raise ValueError(\n            f\"Expected B scales to have a M=128 layout, got {b_scale.layout}\"\n        )\n    else:\n      raise AssertionError(\"Should not happen\")\n    if b_scale.shape[0] % 128 or b_scale.shape[0] < n * num_cta:\n      raise ValueError(\n          f\"B scale shape[0] must be a multiple of 128 and >= N={n * num_cta},\"\n          f\" got {b_scale.shape[0]}\"\n      )\n    if b_scale.shape[1] != k_scales:\n      raise ValueError(\n          f\"B scale shape mismatch: expected ({b_scale.shape[0]}, {k_scales}),\"\n          f\" got {b_scale.shape}\"\n      )\n  if is_sparse:\n    sparse_group_elems = 8 if utils.bitwidth(a_element_type) == 4 else 4\n    # Each sparse group has 2 entries.\n    expected_meta_k = k // sparse_group_elems * 2\n    if a_sparse_metadata.shape != (m, expected_meta_k):\n      raise ValueError(\n          f\"A sparse metadata shape mismatch: expected {(m, expected_meta_k)},\"\n          f\" got {a_sparse_metadata.shape}\"\n      )\n    if a_sparse_metadata.dtype != ir.IntegerType.get_signless(2):","sourceCodeStart":469,"sourceCodeEnd":505,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L469-L505","documentation":"The B scale tensor's leading dimension must be a multiple of 128 and at least n * num_cta to satisfy TMEM allocation granularity for block-scaled MMA. This error reports the offending b_scale.shape[0].","triggerScenarios":"Passing b_scale whose shape[0] is not a multiple of 128 (e.g. 64, 100) or smaller than N (n * num_cta), e.g. N=256 with shape[0]=128.","commonSituations":"Allocating B scales sized exactly to N when N is not a multiple of 128; forgetting to scale up the N dimension for multi-CTA (num_cta=2) kernels.","solutions":["Pad b_scale.shape[0] up to the next multiple of 128 that is >= n * num_cta","Double-check n and num_cta values passed to mma match the shape math"],"exampleFix":"# before\nn, num_cta = 256, 2\nb_scale = make_scales(shape=(128, k_scales))\n# after\nb_scale = make_scales(shape=(512, k_scales))  # multiple of 128 and >= 512","handlingStrategy":"validation","validationCode":"rows = ((n * num_cta + 127) // 128) * 128\nassert b_scale.shape[0] % 128 == 0 and b_scale.shape[0] >= n * num_cta, f'use {rows}'","typeGuard":"def valid_b_scale_rows(rows: int, n: int, num_cta: int) -> bool:\n    return rows % 128 == 0 and rows >= n * num_cta","tryCatchPattern":null,"preventionTips":["Allocate B scales with padding to 128-row granularity","Pass num_cta consistently with the warp-specialization config"],"tags":["gpu","mosaic","tcgen05","shape-mismatch","block-scaling","alignment"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}