{"record":{"id":"bcc79146d72bfcdb","repo":"jax-ml/jax","slug":"async-copies-with-swizzle-require-the-last-dime","errorCode":null,"errorMessage":"Async copies with {swizzle=} require the last dimension of the slice to be exactly {swizzle} bytes i.e.  {(swizzle * 8) // element_bitwidth} elements, but got {slice_shape[-1]} elements.","messagePattern":"Async copies with (.+?) require the last dimension of the slice to be exactly (.+?) bytes i\\.e\\.  (.+?) elements, but got (.+?) elements\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/launch_context.py","lineNumber":1196,"sourceCode":"      del idx  # We overwrote the block index in the loop.\n      if rem_collective_size > 1:\n        raise ValueError(\n            \"None of the leading dimensions in the transformed slice shape\"\n            f\" {slice_shape} is divisible by the collective size\"\n            f\" {collective_size}\"\n        )\n\n    if (zeroth_bw := slice_shape[-1] * element_bitwidth) % 128 != 0:\n      raise ValueError(\n          \"Async copies require the number of bits copied along the last\"\n          f\" dimension to be divisible by 128, but got {zeroth_bw}\"\n      )\n    if (\n        swizzle is not None\n        and swizzle != mgpu_dialect.SwizzlingMode.kNoSwizzle\n        and slice_shape[-1] != (swizzle * 8) // element_bitwidth\n    ):\n      raise ValueError(\n          f\"Async copies with {swizzle=} require the last dimension of the\"\n          f\" slice to be exactly {swizzle} bytes i.e. \"\n          f\" {(swizzle * 8) // element_bitwidth} elements, but got\"\n          f\" {slice_shape[-1]} elements.\"\n      )\n    return (smem_ref, slice_shape, dyn_base_indices, gmem_transform)\n\n  def async_copy(\n      self,\n      *,\n      src_ref: ir.Value,\n      dst_ref: ir.Value,\n      gmem_slice: Any = (),\n      gmem_transform: MemRefTransform | tuple[MemRefTransform, ...] = (),\n      gmem_peer_id: int | ir.Value | GlobalBroadcast | None = None,\n      barrier: utils.BarrierRef | None = None,\n      swizzle: int | None = None,\n      arrive: bool | None = None,","sourceCodeStart":1178,"sourceCodeEnd":1214,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/launch_context.py#L1178-L1214","documentation":"When a swizzle mode (e.g. 32/64/128-byte swizzling) is requested for an async copy, TMA requires the innermost slice dimension to exactly match the swizzle width in bytes: slice_shape[-1] must equal (swizzle * 8) // element_bitwidth elements. Otherwise the copy would produce incorrectly swizzled shared-memory data.","triggerScenarios":"Calling async_copy/async_prefetch with swizzle=SwizzlingMode.k128B (or 32B/64B) and a last-dimension slice whose byte width is not exactly the swizzle width, e.g. 64 elements of f32 with 128B swizzle (256 bytes).","commonSituations":"Configuring SMEM swizzling for tensor-core MMA layouts in Mosaic; changing tile inner size or dtype without updating the swizzle mode to match.","solutions":["Set the innermost slice size to (swizzle_bytes * 8) // element_bitwidth elements (e.g. 32 f32 elements for 128B swizzle).","Or change the swizzle mode to match the existing inner dimension byte width, or use kNoSwizzle/None.","Keep dtype and swizzle consistent when retuning tiles: halving element width requires doubling the element count."],"exampleFix":"// before\nctx.async_copy(..., swizzle=mgpu.SwizzlingMode.k128B, gmem_slice=(slice(0, 64), slice(0, 64)))  # f32 -> 256B inner\n// after\nctx.async_copy(..., swizzle=mgpu.SwizzlingMode.k128B, gmem_slice=(slice(0, 64), slice(0, 32)))  # 32 f32 = 128B","handlingStrategy":"validation","validationCode":"if swizzle is not None and swizzle != mgpu.SwizzlingMode.kNoSwizzle:\n    need = (int(swizzle) * 8) // utils.bitwidth(element_type)\n    assert slice_shape[-1] == need, f'swizzle {swizzle} requires last dim == {need} elements'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Derive swizzle mode from the inner tile byte width rather than hardcoding.","Change swizzle and tile size together when switching dtypes.","Write a helper that picks k32B/k64B/k128B from the last-dim byte size."],"tags":["jax","mosaic-gpu","tma","swizzling","shared-memory"],"backgroundTag":"swizzle-mode-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}