{"record":{"id":"afa7bb50eaeaf19a","repo":"jax-ml/jax","slug":"float64-dot-requires-m-16-n-8-k-16-per-warp-t","errorCode":null,"errorMessage":"float64 dot requires M>=16, N>=8, K>=16 per warp tile (Triton MMAv2 m8n8k4 layout); got {', '.join(errors)}","messagePattern":"float64 dot requires M>=16, N>=8, K>=16 per warp tile \\(Triton MMAv2 m8n8k4 layout\\); got (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/triton/lowering.py","lineNumber":2424,"sourceCode":"  m, k = a_type.shape\n  _, n = b_type.shape\n  if a_type.element_type == ir.F64Type.get():\n    # Triton's MMAv2 fp64 path uses the m8n8k4 PTX instruction but aggregates\n    # it with NumRegisters={m:2, n:1, k:4}, producing an effective m16n8k16\n    # per-warp tile.  Blocks smaller than these minimums cause repM/repN/repK\n    # to round to zero, corrupting the ValueTable and segfaulting the compiler.\n    #   M >= 16  (2 × instrM=8)\n    #   N >=  8  (1 × instrN=8)\n    #   K >= 16  (4 × instrK=4)\n    errors = []\n    if m < 16:\n      errors.append(f\"M={m} < 16\")\n    if n < 8:\n      errors.append(f\"N={n} < 8\")\n    if k < 16:\n      errors.append(f\"K={k} < 16\")\n    if errors:\n      raise ValueError(\n          f\"float64 dot requires M>=16, N>=8, K>=16 per warp tile \"\n          f\"(Triton MMAv2 m8n8k4 layout); got {', '.join(errors)}\"\n      )\n\n  if a_type.element_type != b_type.element_type:\n    raise ValueError(\n        \"a and b must have the same element type, but got:\"\n        f\" {a_type.element_type} and {b_type.element_type}\"\n    )\n\n  assert acc_dtype is not None\n  acc = _zeros(ir.RankedTensorType.get([m, n], _dtype_to_ir_type(acc_dtype)))\n\n  if precision in (\n      lax.DotAlgorithmPreset.BF16_BF16_F32_X3,\n      lax.DotAlgorithmPreset.BF16_BF16_F32_X6,\n      lax.DotAlgorithmPreset.BF16_BF16_F32_X9,\n  ):","sourceCodeStart":2406,"sourceCodeEnd":2442,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/triton/lowering.py#L2406-L2442","documentation":"Triton's fp64 matrix multiply (MMAv2 m8n8k4 PTX path) imposes minimum per-warp tile dimensions of M>=16, N>=8, K>=16. The Pallas lowering validates this for float64 dots and raises ValueError listing which dimensions are too small.","triggerScenarios":"A pallas kernel doing an f64 dot with warp tile shapes like M=8, N=4, or K=8, e.g. tl.dot on float64 blocks of shape (8, 8).","commonSituations":"Tuning block sizes down for memory savings in a double-precision kernel; reusing f32 block sizes (which allow smaller tiles) in an f64 kernel; running on GPUs where fp64 needs the MMAv2 layout.","solutions":["Increase the dot tile so M>=16, N>=8, K>=16 (e.g. use 16x16x16 minimum f64 tiles)","Switch inputs to float32 if full double precision isn't required","Pad the K dimension to a multiple >=16 if the natural K is smaller"],"exampleFix":"# before\nBLOCK_M, BLOCK_N, BLOCK_K = 8, 8, 8  # f64 kernel\nacc = pl.dot(a, b, out_dtype=jnp.float64)\n\n# after\nBLOCK_M, BLOCK_N, BLOCK_K = 16, 16, 16\nacc = pl.dot(a, b, out_dtype=jnp.float64)","handlingStrategy":"validation","validationCode":"def f64_tile_ok(m, n, k):\n    return m >= 16 and n >= 8 and k >= 16\nassert f64_tile_ok(BLOCK_M, BLOCK_N, BLOCK_K)","typeGuard":"def valid_f64_tile(m: int, n: int, k: int) -> bool:\n    return m >= 16 and n >= 8 and k >= 16","tryCatchPattern":null,"preventionTips":["Default f64 kernels to 16x16x16 tiles or larger","Consider float32 unless double precision is strictly required"],"tags":["jax","pallas","triton","float64","tile-shape"],"backgroundTag":"matmul-tile-size-constraint","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}