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

[Yukun-Cui]

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

Summary

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

Testing

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

log2_

dtype	Size	Torch Latency (ms)	Gems Latency (ms)	Speedup	TFLOPS
float16	1073741824	3.380360	3.390120	0.997x	0.317
float16	64, 64	0.005480	0.002520	2.175x	0.002
float16	4096, 4096	0.055680	0.054440	1.023x	0.308
float16	64, 512, 512	0.055760	0.054520	1.023x	0.308
float16	1024, 1024, 1024	3.340440	3.402360	0.982x	0.316
float32	1073741824	6.722140	6.798280	0.989x	0.158
float32	64, 64	0.005440	0.002480	2.194x	0.002
float32	4096, 4096	0.107120	0.103840	1.032x	0.162
float32	64, 512, 512	0.107100	0.104000	1.030x	0.161
float32	1024, 1024, 1024	6.686220	6.781740	0.986x	0.158
bfloat16	1073741824	3.342000	3.408920	0.980x	0.315
bfloat16	64, 64	0.005520	0.002520	2.190x	0.002
bfloat16	4096, 4096	0.055760	0.054200	1.029x	0.310
bfloat16	64, 512, 512	0.055760	0.054280	1.027x	0.309
bfloat16	1024, 1024, 1024	3.359920	3.407560	0.986x	0.315

Operator	Arithmetic Mean Speedup
log2_	1.24x

Files Changed

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

[Dongxu-H]

/test | log2_:mthreads

[Dongxu-H]

/test | log2_:mthreads

[github-actions]

❌ On-demand test failed

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

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

[Dongxu-H]

/test | log2_:mthreads

[Dongxu-H]

LGTM

[Dongxu-H]

/test | log2_:mthreads

[Dongxu-H]

LGTM

[Dongxu-H]

LGTM