[darthvader58]

Description

Fix the lowering of all_to_all_single with uneven split sizes. The new lowering pads each chunk of the input to the input row count, exchanges the chunks with one even stablehlo.all_to_all, and trims each received chunk. An AI assistant (Claude Code) helped to write this change. I reviewed the change.

Related Issue

Fixes #1033

Motivation and Context

The old lowering used an all_gather and slices. It had three problems:

• It built 1-D types. A 2-D input, for example [tokens, hidden], did not compile.
• It took the rows for rank r from the start of the block of each sender. The rows for rank r start after the rows for the ranks below r. Thus, ranks other than rank 0 received incorrect rows.
• It read operand 1 as input_split_sizes and operand 2 as output_split_sizes. The schema has the opposite order.

A rank cannot calculate the correct offsets from the old layout. The offsets need the split sizes of the other ranks, and these are not known at compile time. In the new lowering, the sender uses only its input_split_sizes, and the receiver uses only its output_split_sizes. All types keep the rank of the input. All ranks must have the same input row count, as with the old lowering.

This all-to-all is the token dispatch of expert-parallel MoE models in transformers (ep_dispatch_experts). The issue lists Qwen/Qwen3.6-35B-A3B at EP=8 and Qwen/Qwen3-Coder-Next at EP=16. Even splits do not change. They still lower to one stablehlo.all_to_all.

Type of Change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation update
• Refactoring (no functional changes)

Testing

I updated all-to-all-nonuniform.mlir. It covers a 2-D input with a send matrix that is not symmetric ([[1, 3], [2, 2]], rank 1), a 1-D input, and a chunk with zero rows. Before the fix, the test failed with operand and result must have the same rank.

I built in an arm64 Linux container. WORKSPACE configures LLVM for X86 and NVPTX only, so I added AArch64 to the LLVM targets in a copy of the tree. This change does not affect the conversion pass.

I also ran the pass output for both ranks in the StableHLO interpreter with interpreter.run_parallel and programs = [[@rank0, @rank1]]. Row k of rank i holds 10 * i + k, and the input shape is [4, 8]. Before the fix, the 1-D case of the issue gave [0, 1, 2, 10] on rank 1. After the fix, the 2-D case [[1, 3], [2, 2]] gives the correct rows on both ranks: [0, 10, 11] and [1, 2, 3, 12, 13].

Ran locally (no Neuron hardware needed):

bazel test //third_party/torch-mlir/test/Conversion:TorchToNeuron/all-to-all-nonuniform.mlir.test \
  //third_party/torch-mlir/test/Conversion:TorchToNeuron/all-to-all-single.mlir.test \
  //third_party/torch-mlir/test/Conversion:TorchToNeuron/shard-dim-alltoall.mlir.test
Executed 3 out of 3 tests: 3 tests pass.

stablehlo-translate --interpret run_parallel_fixed.mlir
rank 0: [0, 10, 11]        (3x8, correct)
rank 1: [1, 2, 3, 12, 13]  (5x8, correct)

pre-commit run --all-files
All hooks passed.

Ran on Neuron hardware:

Not run. I do not have a Trainium instance. Please run hardware CI before merge using an actual trainium host.

Requesting Review :
@mgray07 @apoorvtintin

Checklist

• Added or updated tests for this change
• I've performed a self-review for my code
• Ran the tests, or explained under Testing why I could not
• Updated the corresponding documentation
• My changes generate no new warnings
• Said under Description if an AI assistant helped with this change (see AI_POLICY.md)