[torchlib] Support offset-based grouped_mm layouts

[TANGBUDU]

Summary

Builds on the aten_grouped_mm converter added in #2805 and adds support for its offsets form. The existing dense, bias, and cast paths are unchanged.

The converter now handles three layouts:

• 2D x 3D: split rows of the first operand using the offsets.
• 3D x 2D: split columns of the second operand using the offsets.
• 2D x 2D: split the contraction dimension and stack the group results.

The number of groups must be known at export time, but offset values and input lengths can remain dynamic. If the final offset does not cover the full input, the converter keeps PyTorch's full output shape and fills the unwritten tail with zeros.

The offsets path does not currently support bias or changing the output dtype. The existing dense converter still supports both.

Testing

OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 MKL_NUM_THREADS=1 \
python -m pytest \
  tests/function_libs/torch_lib/grouped_mm_test.py \
  tests/function_libs/torch_lib/ops_test.py \
  -k grouped_mm -q

• PyTorch 2.14 environment: 50 passed, 1 skipped.
• Repository-pinned environment: 46 passed, 5 skipped.
• Tests cover all three layouts, uneven and empty groups, dynamic offsets, dynamic input lengths, and a single group.
• Three dynamic BF16 exports pass the ONNX checker. BF16 runtime and GPU execution were not tested.
• lintrunner -a and git diff --check pass.

AI assistance

ChatGPT assisted with the implementation, tests, and result analysis. I reviewed the changes and the test results above.

[TANGBUDU]

@microsoft-github-policy-service agree