{"record":{"id":"ea252e0136a448dd","repo":"jax-ml/jax","slug":"no-tuned-tiling-found-for-m-k-n-m-k","errorCode":null,"errorMessage":"No tuned tiling found for (m, k, n) = ({m}, {k}, {n})","messagePattern":"No tuned tiling found for \\(m, k, n\\) = \\((.+?), (.+?), (.+?)\\)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/megablox/gmm.py","lineNumber":376,"sourceCode":"    group_offset = group_offset[None]\n  num_current_groups = rhs.shape[0]\n  num_total_groups = group_sizes.shape[0]\n  lhs, group_sizes, input_dtype = _validate_args(\n      lhs=lhs, rhs=rhs, group_sizes=group_sizes\n  )\n\n  # Gather shape information.\n  m, k, n = (lhs.shape[0], lhs.shape[1], rhs.shape[2])\n  if transpose_rhs:\n    n = rhs.shape[1]\n\n  # If tiling is callable, look up the problem dimensions in the LUT. If no tuned\n  # tile dimensions are available throw an error.\n  if callable(tiling):\n    tiling = tiling(m, k, n)\n\n  if tiling is None:\n    raise ValueError(f\"No tuned tiling found for (m, k, n) = ({m}, {k}, {n})\")\n\n  tm, tk, tn = tiling\n  tiles_k, k_rem = _calculate_irregular_num_tiles(k, tk)\n  tiles_n, n_rem = _calculate_irregular_num_tiles(n, tn)\n  del n_rem\n\n  # Create the metadata we need for computation.\n  group_metadata, num_active_tiles = make_group_metadata(\n      group_sizes=group_sizes,\n      m=m,\n      tm=tm,\n      start_group=group_offset[0],\n      num_nonzero_groups=rhs.shape[0],\n      visit_empty_groups=False,\n  )\n\n  def kernel(\n      group_metadata,","sourceCodeStart":358,"sourceCodeEnd":394,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/megablox/gmm.py#L358-L394","documentation":"Megablox gmm selects tile sizes from a tuned lookup table (LUT) keyed by (m, k, n) problem dimensions. When the tiling argument is a callable LUT and it returns None for the given dimensions, no tuned configuration exists and the kernel refuses to guess, raising this ValueError.","triggerScenarios":"Calling gmm with default tiling=select_tiling (or another callable) for a matrix shape absent from the tuning LUT, e.g. unusual m=1333 or k dimensions not covered by tuning runs. Also raised by _gmm_fwd/_gmm_bwd during autodiff of such a call.","commonSituations":"Changing batch size or model dims to a value never tuned (e.g. non-multiples of 128); using a newer/older JAX whose LUT covers a different shape set; running MoE experts with atypical shapes.","solutions":["Pass an explicit tiling tuple, e.g. tiling=(128, 128, 128), instead of the LUT callable","Round m/k/n to dimensions present in the tuned table (typically multiples of 128)","Use jnp.einsum/dot as a fallback for untuned shapes"],"exampleFix":"// before\nout = gmm(lhs, rhs, group_sizes)  # default LUT\n// after\nout = gmm(lhs, rhs, group_sizes, tiling=(128, 128, 128))","handlingStrategy":"fallback","validationCode":"tiling = select_tiling(m, k, n)\nif tiling is None:\n    tiling = (128, 128, 128)  # safe default for your workload","typeGuard":null,"tryCatchPattern":"try:\n    out = gmm(lhs, rhs, group_sizes)\nexcept ValueError as e:\n    if 'No tuned tiling' in str(e):\n        out = gmm(lhs, rhs, group_sizes, tiling=(128, 128, 128))\n    else:\n        raise","preventionTips":["Probe select_tiling(m, k, n) first and cache explicit tilings per shape","Keep problem dims on tuned multiples of 128","Wrap megablox calls with a dot-product fallback for benchmarking safety"],"tags":["jax","pallas","tpu","megablox","tiling","unsupported-shape"],"backgroundTag":"unsupported-configuration","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}