Skip to content

Add mm operator support - #38

Merged
voltjia merged 3 commits into
masterfrom
operators/mm
Jul 28, 2026
Merged

Add mm operator support#38
voltjia merged 3 commits into
masterfrom
operators/mm

Conversation

@voltjia

@voltjia voltjia commented Jul 27, 2026

Copy link
Copy Markdown
Collaborator

Summary

This PR implements aten::mm for two-dimensional, same-device float32 Infini tensors by calling InfiniOps Gemm implementation index 0 on the current stream.

The initial path supports contiguous inputs, transposed non-overlapping dense weights, empty dimensions, and input/output storage lifetime tracking. It intentionally rejects unsupported dtypes, devices, ranks, and overlapping layouts instead of falling back through CPU.

This enables torch.nn.Linear(..., bias=False) inference through PyTorch's existing transpose-and-mm path. Bias, mm.out, training, backward, and propagation of torch.set_float32_matmul_precision("highest") remain outside this PR.

This PR is stacked on #37, which is stacked on #36. The intended merge order is #36, #37, and then this PR.

Testing

  • Ruff 0.15.21 and clang-format 19.1.4 pass.
  • An installed wheel built against PyTorch 2.12.0+cpu passes the focused suite with 11 passed, 1 skipped, and 1 xfailed, and the full suite with 181 passed, 43 skipped, and 3 xfailed.
  • An installed wheel built against PyTorch 2.13.0+cpu passes the focused suite with 11 passed, 1 skipped, and 1 xfailed, and the full suite with 181 passed, 43 skipped, and 3 xfailed.
  • An installed wheel built against PyTorch 2.13.0+cu130 passes the focused suite on an NVIDIA A100-SXM4-80GB with 11 passed, 1 skipped, and 1 xfailed.
  • The NVIDIA validation uses driver 580.105.08, CUDA toolkit 13.1.80, CUDA architecture 80, InfiniRT commit 95c70080f9551e61241110497d163dfcdf9dc7e7, and InfiniOps commit 296271487beb594a248fd463e5fff14f7ab74293.
  • The Infini results for contiguous and transposed-weight inputs match official torch.cuda.mm under PyTorch 2.13's default high float32 matrix-multiplication precision.
  • The acceptance suite runs an actual bias-free torch.nn.Linear module after moving its weight and input to the infini device.

Base automatically changed from metadata/as-strided to master July 28, 2026 02:22
@voltjia
voltjia merged commit 8aee853 into master Jul 28, 2026
8 checks passed
@voltjia
voltjia deleted the operators/mm branch July 28, 2026 02:23
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant