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

[Yukun-Cui]

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

Summary

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

Testing

• Reused the existing upstream accuracy tests tests/test_fix.py (-m fix)
• Validated against reference on the MUSA device; specialization confirmed active via the GEMS_MTHREADS FIX 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).

fix

dtype	Size	Torch Latency (ms)	Gems Latency (ms)	Speedup	TFLOPS
float16	1073741824	3.241560	3.106380	1.044x	0.346
float16	64, 64	0.004960	0.002640	1.879x	0.002
float16	4096, 4096	0.078560	0.078000	1.007x	0.215
float16	64, 512, 512	0.078080	0.077600	1.006x	0.216
float16	1024, 1024, 1024	3.221280	3.106060	1.037x	0.346
float32	1073741824	6.401800	6.412680	0.998x	0.167
float32	64, 64	0.004860	0.002640	1.841x	0.002
float32	4096, 4096	0.132640	0.131280	1.010x	0.128
float32	64, 512, 512	0.133760	0.131880	1.014x	0.127
float32	1024, 1024, 1024	6.394120	6.408480	0.998x	0.168
bfloat16	1073741824	3.116800	3.103780	1.004x	0.346
bfloat16	64, 64	0.004780	0.002720	1.757x	0.002
bfloat16	4096, 4096	0.079040	0.078000	1.013x	0.215
bfloat16	64, 512, 512	0.078560	0.077540	1.013x	0.216
bfloat16	1024, 1024, 1024	3.114980	3.104640	1.003x	0.346

Operator	Arithmetic Mean Speedup
fix	1.17x

Files Changed

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

[Dongxu-H]

/test | fix:mthreads

[github-actions]

❌ On-demand test failed

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

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

[Dongxu-H]

/test | fix:mthreads

[Dongxu-H]

/test | fix:mthreads

[Dongxu-H]

LGTM

[Dongxu-H]

/test | fix:mthreads