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

[Yukun-Cui]

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

Summary

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

Testing

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

channel_shuffle

dtype	Size	Torch Latency (ms)	Gems Latency (ms)	Speedup
float16	1, 4, 2, 2	0.004200	0.004320	0.972x
float16	2, 8, 4, 4	0.004440	0.004320	1.028x
float16	4, 16, 8, 8	0.004480	0.004400	1.018x
float32	1, 4, 2, 2	0.004200	0.004440	0.946x
float32	2, 8, 4, 4	0.004360	0.004400	0.991x
float32	4, 16, 8, 8	0.004640	0.005520	0.841x
bfloat16	1, 4, 2, 2	0.004200	0.004480	0.938x
bfloat16	2, 8, 4, 4	0.004400	0.004520	0.973x
bfloat16	4, 16, 8, 8	0.004480	0.004560	0.982x

Operator	Arithmetic Mean Speedup
channel_shuffle	0.97x

Files Changed

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

[Dongxu-H]

/test | channel_shuffle:mthreads

[github-actions]

❌ On-demand test failed

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

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

[Dongxu-H]

/test | channel_shuffle:mthreads

[Dongxu-H]

LGTM

[Dongxu-H]

LGTM