{"record":{"id":"e489aaad9bc7798d","repo":"jax-ml/jax","slug":"group-offset-is-not-currently-supported-in-the-pal","errorCode":null,"errorMessage":"group_offset is not currently supported in the pallas-triton lowering.","messagePattern":"group_offset is not currently supported in the pallas-triton lowering\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/pallas_lowerings/gpu/ragged_dot.py","lineNumber":467,"sourceCode":"\ndef _pallas_ragged_dot_general_impl(\n    lhs: Array,\n    rhs: Array,\n    group_sizes: Array,\n    ragged_dot_dimension_numbers: lax.RaggedDotDimensionNumbers,\n    precision: lax.PrecisionLike = None,  # TODO(rdyro): Add support for the\n                                          # precision parameter in the kernels.\n    preferred_element_type: DTypeLike | None = None,\n    group_offset: Array | None = None,\n    out_sharding=None,\n) -> Array:\n  if out_sharding is not None:\n    raise NotImplementedError(\n      \"Explicit sharding is not currently supported in the pallas-triton\"\n      \" lowering of ragged_dot_general. You can call this op under shard_map.\")\n\n  if group_offset is not None:\n    raise NotImplementedError(\"group_offset is not currently supported in the \"\n                              \"pallas-triton lowering.\")\n\n  mode, lhs_ragged_dim = lax._ragged_dot_mode_and_dim(\n      lhs.ndim, ragged_dot_dimension_numbers)\n  (l_contract, r_contract), (l_batch, r_batch) = (\n      ragged_dot_dimension_numbers.dot_dimension_numbers\n  )\n  l_contract, r_contract = tuple(l_contract), tuple(r_contract)\n  l_batch, r_batch = tuple(l_batch), tuple(r_batch)\n  l_noncontract = tuple(lax.remaining(range(lhs.ndim), l_contract, l_batch))\n  assert len(l_batch) == len(r_batch)\n\n  if group_sizes.ndim == 1:\n    group_sizes = lax.broadcast(group_sizes, [lhs.shape[i] for i in l_batch])\n\n  compute_dtype = np.promote_types(lhs.dtype, rhs.dtype)\n  out_dtype = (np.dtype(preferred_element_type)\n               if preferred_element_type is not None else None)","sourceCodeStart":449,"sourceCodeEnd":485,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/pallas_lowerings/gpu/ragged_dot.py#L449-L485","documentation":"The pallas-triton lowering of ragged_dot_general does not support the group_offset argument, which allows applying only a sub-range of groups. Only full grouped dot semantics are implemented on this backend.","triggerScenarios":"Passing a non-None group_offset to ragged_dot_general when lowering through pallas-triton on GPU (e.g. chunked/pipelined MoE group processing).","commonSituations":"Reusing TPU MoE code that slices expert groups via group_offset; attempting incremental grouped matmul to cap memory on GPU.","solutions":["Remove group_offset and process all groups at once (slice the weights/activations manually instead)","Pre-slice A and group_sizes so the desired group subrange becomes the full call","Use the XLA ragged_dot_general path rather than the pallas-triton lowering if group_offset is required"],"exampleFix":"// before\nout = ragged_dot_general(lhs, rhs, dims, group_sizes, group_offset=off)\n\n// after\nA_sub, gs_sub = A[off:off+n], group_sizes[off:off+n]\nout = ragged_dot_general(lhs, A_sub, dims, gs_sub)","handlingStrategy":"fallback","validationCode":"if group_offset is not None and backend == 'triton':\n    raise SystemExit('slice A/group_sizes manually instead of group_offset')","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Avoid backend-exclusive args in portable kernels","Slice weights before calling ragged_dot_general"],"tags":["jax","pallas","gpu","not-implemented","ragged-dot"],"backgroundTag":"feature-not-implemented","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}