{"record":{"id":"270253114d51a17e","repo":"jax-ml/jax","slug":"wgmma-requires-m-and-n-to-be-multiples-of-64-and-8","errorCode":null,"errorMessage":"WGMMA requires m and n to be multiples of 64 and 8, got {m} and {n}","messagePattern":"WGMMA requires m and n to be multiples of 64 and 8, got (.+?) and (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/wgmma.py","lineNumber":69,"sourceCode":"      self,\n      *,\n      _value: fa.FragmentedArray,\n      _original_layout: fa.FragmentedLayout,\n      _sync: bool = True,\n  ):\n    self._original_layout = _original_layout\n    self._value = _value\n    if _sync:\n      self._value = wgmma_fence(_value)\n\n  @property\n  def value(self) -> fa.FragmentedArray:\n    return self._value.to_layout(self._original_layout)\n\n  @classmethod\n  def zero(cls, m, n, dtype=None, *, is_signed: bool | None = None):\n    if m % 64 or n % 8:\n      raise ValueError(\"WGMMA requires m and n to be multiples of 64 and 8, \"\n                       f\"got {m} and {n}\")\n    if is_signed is False:\n      raise TypeError(\"PTX does not support unsigned WGMMA accumulators\")\n    f32 = ir.F32Type.get()\n    if dtype is None:\n      dtype = f32\n    if isinstance(dtype, ir.IntegerType):\n      zero = arith.constant(dtype, ir.IntegerAttr.get(dtype, 0))\n    else:\n      zero = arith.constant(dtype, ir.FloatAttr.get(dtype, 0.0))\n    return cls.from_registers(\n        fa.FragmentedArray.splat(\n            zero, (m, n), fa.WGMMA_LAYOUT, is_signed=is_signed\n        )\n    )\n\n  @classmethod\n  def from_registers(cls, registers, sync=True):","sourceCodeStart":51,"sourceCodeEnd":87,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/wgmma.py#L51-L87","documentation":"WGMMAAccumulator.zero (wgmma.py:69) validates WGMMA tile geometry: the Hopper wgmma.mma_async instruction requires the M dimension to be a multiple of 64 and N a multiple of 8. Non-conforming shapes cannot be executed by the tensor cores.","triggerScenarios":"Calling wgmma.WGMMAAccumulator.zero(m, n) with m not divisible by 64 or n not divisible by 8 (e.g. zero(64, 12) or zero(32, 16)).","commonSituations":"Padding attention/GEMM tiles incorrectly; deriving m/n from head_dim or batch sizes without rounding up; porting TMA layouts whose tile shape is not WGMMA-conformant.","solutions":["Round m up to a multiple of 64 and n up to a multiple of 8 (padding as needed)","Restructure the kernel loop so each wgmma call operates on conformant tiles","Add a shape assertion early so the failure is caught at kernel-definition time"],"exampleFix":"# before\nacc = wgmma.WGMMAAccumulator.zero(m=48, n=128)\n# after\nm = (m + 63) // 64 * 64\nn = (n + 7) // 8 * 8\nacc = wgmma.WGMMAAccumulator.zero(m=m, n=n)","handlingStrategy":"validation","validationCode":"assert m % 64 == 0 and n % 8 == 0, f'WGMMA tiles must satisfy m%64==0, n%8==0 (got {m}, {n})'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Define tile sizes as multiples of (64, 8) constants","Pad dynamic shapes up front in kernel launchers"],"tags":["jax","mosaic-gpu","wgmma","gpu-kernel","shape-validation"],"backgroundTag":"gpu-tile-shape-invalid","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}