【KernelGen】Add einsum operator

[Schopenhauer-loves-Hegel]

PR Category

Operator

Type of Change

New Feature

Description

Add einsum operator implementation with Triton kernel.

• Implementation mode: N/A
• Accuracy test: 57/57 passed

Issue

N/A

Progress

• Change is properly reviewed (1 reviewer required, 2 recommended).
• Change is responded to an issue.
• Change is fully covered by a UT.

Performance

torch.bfloat16

Shape	Torch Latency (ms)	Gems Latency (ms)	Speedup
torch.Size([384, 384])	0.0109	0.0119	0.912
torch.Size([4096, 4096])	0.5583	0.5640	0.990
torch.Size([1024, 1024])	0.0229	0.0263	0.872
torch.Size([2048, 2048])	0.1116	0.0996	1.121
torch.Size([4096, 4096])	0.5459	0.5637	0.968
torch.Size([2, 384, 384])	0.0130	0.0130	1.000
torch.Size([2, 4096, 4096])	1.0824	1.1406	0.949
torch.Size([16, 1024, 1024])	0.1536	0.1606	0.957
torch.Size([16, 2048, 2048])	1.1166	1.1304	0.988
torch.Size([16, 4096, 4096])	8.8083	9.2300	0.954

torch.float16

Shape	Torch Latency (ms)	Gems Latency (ms)	Speedup
torch.Size([384, 384])	0.0140	0.0122	1.144
torch.Size([4096, 4096])	0.5663	0.5762	0.983
torch.Size([1024, 1024])	0.0209	0.0258	0.811
torch.Size([2048, 2048])	0.0787	0.0988	0.796
torch.Size([4096, 4096])	0.5729	0.5715	1.003
torch.Size([2, 384, 384])	0.0130	0.0135	0.957
torch.Size([2, 4096, 4096])	1.1233	1.1739	0.957
torch.Size([16, 1024, 1024])	0.1568	0.1627	0.964
torch.Size([16, 2048, 2048])	1.1334	1.1718	0.967
torch.Size([16, 4096, 4096])	9.0352	9.4699	0.954

torch.float32

Shape	Torch Latency (ms)	Gems Latency (ms)	Speedup
torch.Size([384, 384])	0.0225	0.0340	0.661
torch.Size([4096, 4096])	7.2281	8.0535	0.898
torch.Size([1024, 1024])	0.1361	0.1938	0.702
torch.Size([2048, 2048])	0.9777	1.1970	0.817
torch.Size([4096, 4096])	7.2291	8.0542	0.898
torch.Size([2, 384, 384])	0.0827	0.0263	3.144
torch.Size([2, 4096, 4096])	14.3647	14.5997	0.984
torch.Size([16, 1024, 1024])	1.9076	1.8544	1.029
torch.Size([16, 2048, 2048])	14.3852	14.6119	0.984
torch.Size([16, 4096, 4096])	114.7864	116.1884	0.988

Overall: median speedup = 0.960x, mean speedup = 1.012x (30 data points)

---

Generated by auto_gen tool with Claude Code

[Schopenhauer-loves-Hegel]

Closing in favor of grouped PRs to reduce review overhead.

[CLAassistant]

All committers have signed the CLA.

[factnn]

einsum Benchmark Speedup Results

einsum_matmul (float16)

Size	Torch (ms)	Gems (ms)	Speedup
[384,384]x[384,384]	0.013	0.015	0.866
[1024,1024]x[1024,1024]	0.022	0.027	0.813
[2048,2048]x[2048,2048]	0.077	0.104	0.740
[4096,4096]x[4096,4096]	0.545	0.748	0.728

einsum_matmul (float32)

Size	Torch (ms)	Gems (ms)	Speedup
[384,384]x[384,384]	0.023	0.043	0.528
[1024,1024]x[1024,1024]	0.135	0.193	0.703
[4096,4096]x[4096,4096]	7.225	8.063	0.896

einsum_bmm (float16)

Size	Torch (ms)	Gems (ms)	Speedup
[2,384,384]x2	0.013	0.015	0.860
[16,1024,1024]x2	0.155	0.174	0.890
[16,2048,2048]x2	1.135	1.247	0.911
[16,4096,4096]x2	8.872	10.185	0.871

einsum_bmm (float32)

Size	Torch (ms)	Gems (ms)	Speedup
[2,384,384]x2	0.082	0.030	2.764
[16,1024,1024]x2	1.905	1.860	1.024
[16,4096,4096]x2	114.759	116.535	0.985

[tengqm]

Some naming suggestions and sorting nits.
Wondering if we are doing a thorough performance testing when compared to accuracy checking. There are so many situations to cover.

[factnn]

Thanks for the review! All three comments have been addressed:

1. EINSUM_SHAPES style changed to if QUICK_MODE: ... else: ... pattern
2. Removed duplicate standalone einsum_bmm_op / einsum_matmul_op functions, inlined as lambdas
3. Regarding reusing existing shapes: einsum covers multiple operation patterns (matmul, bmm, dot, outer, trace, transpose, sum), each requiring different shape dimensions. The existing POINTWISE_SHAPES, REDUCTION_SHAPES, MNK_SHAPES etc. don't directly fit these varied formats, so custom shapes are kept for clarity.

[tengqm]

overall lgtm. thanks.

[factnn]

UT Results

Branch rebased to upstream/master and tests passed:

python -m pytest tests/ -m einsum --ref cpu -x -q
60 passed

[factnn]

The ii->i pattern is supported — it's handled by _is_diagonal_pattern before reaching _is_sum_pattern, so the uniqueness check in _is_sum_pattern doesn't affect it. Verified:

A = torch.randn(3, 3, device='cuda')
torch.einsum('ii->i', A)  # matches FlagGems output exactly

[factnn]

The ii->i pattern is supported — it's handled by _is_diagonal_pattern before reaching _is_sum_pattern, so the uniqueness check in _is_sum_pattern doesn't affect it. Verified locally and results match exactly.