{"record":{"id":"ed02b14eedb91514","repo":"jax-ml/jax","slug":"async-copy-requires-all-gmem-strides-except-the-la","errorCode":null,"errorMessage":"async_copy requires all GMEM strides except the last one to be a multiple of 16 bytes","messagePattern":"async_copy requires all GMEM strides except the last one to be a multiple of 16 bytes","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/launch_context.py","lineNumber":1075,"sourceCode":"      swizzle: int | None,\n      slice_shape: list[int],\n      dyn_base_indices: tuple[ir.Value, ...],\n      gather_indices,\n      squeezed_dims: tuple[int, ...],\n      gmem_transform: tuple[MemRefTransform, ...],\n      collective: Sequence[gpu.Dimension],\n      leader_tracked: CopyPartition | None = None,\n  ):\n    \"\"\"Finalizes setup specific to the TMA implementation of async_copy.\"\"\"\n    index = ir.IndexType.get()\n    # The function below is called only to verify the GMEM ref. The output\n    # is meant to be ignored.\n    _find_kernel_argument_for_gmem_ref(gmem_ref)\n    gmem_ref_ty = ir.MemRefType(gmem_ref.type)\n    element_bitwidth = utils.bitwidth(gmem_ref_ty.element_type)\n    gmem_strides, _ = gmem_ref_ty.get_strides_and_offset()\n    if any(s * element_bitwidth % 128 != 0 for s in gmem_strides[:-1]):\n      raise ValueError(\n          \"async_copy requires all GMEM strides except the last one to be a\"\n          \" multiple of 16 bytes\"\n      )\n    # We don't need to do this for gather TMAs, because we'll unroll the\n    # transfers ourselves anyway.\n    num_squeezed_dims = len(squeezed_dims)\n    if gather_indices is None:\n      # Drop as many unit-sized dimensions from the transformed shape as we can.\n      gmem_shape = tuple(gmem_ref_ty.shape)\n      for t in gmem_transform:\n        gmem_shape = t.transform_gmem_shape(gmem_shape)\n      # The slice shape may pad along 1-sized dimensions. In that case, we do\n      # not drop them.\n      unit_dims = tuple(\n          i for i, (gs, ss) in enumerate(zip(gmem_shape, slice_shape, strict=True))\n          if gs == 1 and ss == 1\n      )\n      # When issuing an `async_prefetch`, there is no SMEM reference to","sourceCodeStart":1057,"sourceCodeEnd":1093,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/launch_context.py#L1057-L1093","documentation":"The TMA (Tensor Memory Accelerator) hardware path requires that global-memory tensor strides, except the innermost one, be multiples of 16 bytes. The code checks each stride (in elements) times the element bitwidth is divisible by 128 bits; any outer stride violating this raises the error. This mirrors NVIDIA TMA's global-memory alignment requirement.","triggerScenarios":"Calling async_copy or async_prefetch with implementation=AsyncCopyImplementation.TMA on a GMEM reference whose outer strides * element_bitwidth are not multiples of 128 bits, e.g. an f32 tensor with a row stride of 3 elements (12 bytes).","commonSituations":"Passing non-contiguous or oddly strided views into a Mosaic kernel; using small leading dimensions (e.g. shape (3, N) f32) in TMA copies; switching a kernel from the non-TMA path to TMA on Hopper+ GPUs.","solutions":["Pad or reallocate the global tensor so all outer strides are multiples of 16 bytes (e.g. pad leading dims so stride * element_size % 16 == 0).","Check and canonicalize the input layout before the kernel (jnp.reshape/pad to a contiguous 16-byte-aligned layout).","If alignment is impossible, use a non-TMA AsyncCopyImplementation.","Verify utils.bitwidth assumptions: for sub-32-bit types the stride in elements must be proportionally larger."],"exampleFix":"// before\nx = jnp.zeros((3, 128), dtype=jnp.float32)  # row stride 3*4=12 bytes\nctx.async_copy(gmem_ref, smem_ref, ..., implementation=mgpu.AsyncCopyImplementation.TMA)\n// after\nx = jnp.zeros((4, 128), dtype=jnp.float32)  # padded so stride is 16-byte aligned\nctx.async_copy(gmem_ref, smem_ref, ..., implementation=mgpu.AsyncCopyImplementation.TMA)","handlingStrategy":"validation","validationCode":"et = gmem_ref.type.element_type\nbw = utils.bitwidth(et)\nstrides, _ = ir.MemRefType(gmem_ref.type).get_strides_and_offset()\nassert all(s * bw % 128 == 0 for s in strides[:-1]), 'outer strides must be 16-byte aligned for TMA'","typeGuard":null,"tryCatchPattern":"try:\n    ctx.async_copy(..., implementation=mgpu.AsyncCopyImplementation.TMA)\nexcept ValueError as e:\n    if 'GMEM strides' in str(e):\n        x = jnp.ascontiguousarray(_pad_to_16B_align(x))\n    else:\n        raise","preventionTips":["Pad input tensors so leading dims give 16-byte-aligned row strides.","Prefer contiguous layouts for tensors passed to TMA kernels.","Add a stride-alignment assert in kernel prologue code."],"tags":["jax","mosaic-gpu","tma","memory-alignment","strides"],"backgroundTag":"memory-alignment-error","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}