{"record":{"id":"9206e33d86278e5e","repo":"jax-ml/jax","slug":"cannot-assign-layout-to-async-load-with-gather-ind","errorCode":null,"errorMessage":"Cannot assign layout to async load with gather indices since minor dim={slice_lengths[-1]} is not divisible by {divisor=} bits.","messagePattern":"Cannot assign layout to async load with gather indices since minor dim=(.+?) is not divisible by (.+?) bits\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/layout_inference.py","lineNumber":2247,"sourceCode":"      tiling_multiple.append(size)\n      continue\n    tiling_multiple.append(dynamic_gcd(size, index))\n\n  operand_index = 1 if isinstance(op, mgpu.AsyncLoadOp) else 0\n  operand = ValueSite(op, VariableType.OPERAND, operand_index)\n  var = ctx.producer_ref(operand)\n  constraints: list[cs.Constraint] = [\n      cs.Divides(expr=var, tiling_multiple=tuple(tiling_multiple))\n  ]\n  if any(isinstance(idx.type, ir.VectorType) for idx in op.indices):\n    element_bitwidth = utils.bitwidth(op.source.type.element_type)\n    # This constraint enforces sufficient SMEM-alignment.\n    # The transfer chunk needs to be 1024 bit-aligned. For each write in the\n    # lowering we transfer 4 rows, so each row must be 256 bit-aligned.\n    divisor = (1024 // 4) // element_bitwidth\n    slice_lengths = [s for s in op.slice_lengths if s != -1]\n    if slice_lengths and (slice_lengths[-1] % divisor):\n      raise ValueError(\n          \"Cannot assign layout to async load with gather indices since\"\n          f\" minor dim={slice_lengths[-1]} is not divisible by {divisor=}\"\n          \" bits.\"\n      )\n    constraints.append(cs.MinorDimDivisibleBy(expr=var, divisor=divisor))\n\n  value_sites_for_variable = {var: [operand]}\n  value_sites, extra_constraints = _vector_value_sites_and_constraints_for_async_ops(op)\n  value_sites_for_variable.update(value_sites)\n  constraints.extend(extra_constraints)\n  return cs.ConstraintSystem(constraints=constraints), value_sites_for_variable\n\n\n@_add_constraint_system_derivation_rule(mgpu.AsyncPrefetchOp)\ndef _async_prefetch_constraint_system(\n    ctx: DerivationContext,\n    op: mgpu.AsyncPrefetchOp,\n) -> ConstraintSystemDerivationRuleResult:","sourceCodeStart":2229,"sourceCodeEnd":2265,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/layout_inference.py#L2229-L2265","documentation":"Mosaic GPU layout inference enforces SMEM alignment for async loads/stores with gather indices: each transferred row must be 256-bit aligned, so the minor dimension length of the slice must be divisible by (256 // element_bitwidth). If the last non-(-1) slice length isn't divisible by that divisor, no valid layout can be assigned and a ValueError is raised.","triggerScenarios":"Calling mgpu.async_load (or async store) with gather/slice indices whose innermost slice length (excluding -1 dims) is not a multiple of 256//bitwidth, e.g. minor dim 4 with 16-bit elements (divisor 16).","commonSituations":"Hand-written Mosaic kernels doing TMA/tensor gather loads with narrow inner dimensions (fp16/bf16 with small tiles), or padding vectors to non-multiple-of-16 element counts.","solutions":["Pad the innermost (minor) dimension of the gathered slice so its length is divisible by 256//element_bitwidth (e.g. multiples of 16 for 16-bit types)","Use a wider element type or reshape so the minor dim satisfies the alignment","Avoid gather indices: use a regular async load which doesn't hit this constraint"],"exampleFix":"// before (bf16, minor dim 8 -> divisor 16)\nvals = mgpu.async_load(..., slice_lengths=(..., 8), ...)\n// after\nvals = mgpu.async_load(..., slice_lengths=(..., 16), ...)  # pad minor dim to 16","handlingStrategy":"validation","validationCode":"bitwidth = 16  # e.g. bf16\ndivisor = 256 // bitwidth\nassert slice_lengths[-1] % divisor == 0, f\"minor dim {slice_lengths[-1]} not divisible by {divisor}\"","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always size the innermost gathered dimension as a multiple of 16 elements for 16-bit types","Compute 256//bitwidth and validate before issuing async_load with gather indices"],"tags":["jax","mosaic","gpu","layout","alignment","async-load"],"backgroundTag":"memory-alignment-constraint-violation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}