[KernelGen][MThreads] Add square_ Moore Threads specialized operator

[Yukun-Cui]

[KernelGen][MThreads] Add square_ Moore Threads specialized operator

Summary

Add a Moore Threads (MUSA) specialized Triton kernel for square_, overriding the generic implementation via runtime.replace_customized_ops(). Implements the square operation.

Testing

• Reused the existing upstream accuracy tests tests/test_square.py (-m square_)
• Validated against reference on the MUSA device; specialization confirmed active via the GEMS_MTHREADS SQUARE_ debug log
• Falls back to the generic implementation for unsupported dtype/device/shape (fp64/int64 are not supported on Moore Threads hardware)

Performance

Compared against the generic FlagGems implementation on Moore Threads (MUSA).

square_

dtype	Size	Torch Latency (ms)	Gems Latency (ms)	Speedup	TFLOPS
float16	1073741824	3.440360	3.538760	0.972x	0.303
float16	64, 64	0.011000	0.003840	2.865x	0.001
float16	4096, 4096	0.054880	0.054080	1.015x	0.310
float16	64, 512, 512	0.055120	0.054160	1.018x	0.310
float16	1024, 1024, 1024	3.395960	3.552040	0.956x	0.302
float32	1073741824	6.706220	6.868860	0.976x	0.156
float32	64, 64	0.007960	0.003160	2.519x	0.001
float32	4096, 4096	0.105920	0.104560	1.013x	0.160
float32	64, 512, 512	0.105840	0.104520	1.013x	0.161
float32	1024, 1024, 1024	6.732220	6.913800	0.974x	0.155
bfloat16	1073741824	3.405880	3.525880	0.966x	0.305
bfloat16	64, 64	0.010800	0.003640	2.967x	0.001
bfloat16	4096, 4096	0.054960	0.054200	1.014x	0.310
bfloat16	64, 512, 512	0.055120	0.054360	1.014x	0.309
bfloat16	1024, 1024, 1024	3.442680	3.514580	0.980x	0.306

Operator	Arithmetic Mean Speedup
square_	1.35x

Files Changed

• src/flag_gems/runtime/backend/_mthreads/ops/square_.py: Moore Threads Triton kernel + fallback
• src/flag_gems/runtime/backend/_mthreads/ops/__init__.py: Register import and __all__

[Dongxu-H]

/test | square_:mthreads

[github-actions]

❌ On-demand test failed

Operator: | square_
Runner: mthreads
Backend: mthreads-musa520

The test failed to complete. Check the workflow run for details.

[Dongxu-H]

LGTM

[Dongxu-H]

This PR contains 142 files with large infrastructure changes (CI workflows, CODEOWNERS deletions, benchmark modifications) mixed with the square_ operator implementation. This makes focused code review difficult. Consider splitting into separate PRs for operator implementation versus infrastructure changes. The square_ implementation in mthreads backend should be reviewed separately. Mthreads backend CI passes but the PR needs focused review on the operator-specific code which is buried in 142 files.

[Dongxu-H]

LGTM