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

[Yukun-Cui]

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

Summary

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

Testing

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

reflection_pad3d_backward

dtype	Size	Torch Latency (ms)	Gems Latency (ms)	Speedup
float16	1, 1, 6, 6, 6	0.020960	0.010880	1.926x
float16	1, 1, 8, 8, 8	0.037940	0.011560	3.282x
float16	1, 1, 7, 7, 7	0.032040	0.011120	2.881x
float16	2, 3, 10, 10, 10	0.017840	0.010920	1.634x
float16	2, 3, 12, 12, 12	0.022160	0.010960	2.022x
float16	2, 3, 11, 11, 11	0.021060	0.010880	1.936x
float16	1, 1, 18, 18, 18	0.018560	0.007680	2.417x
float16	1, 1, 20, 20, 20	0.025040	0.007760	3.227x
float16	1, 1, 19, 19, 19	0.024480	0.007760	3.155x
float16	2, 4, 10, 18, 34	0.020560	0.007840	2.622x
float16	2, 4, 12, 20, 36	0.026200	0.008240	3.180x
float16	2, 4, 11, 19, 35	0.022680	0.008120	2.793x
float32	1, 1, 6, 6, 6	0.014040	0.010240	1.371x
float32	1, 1, 8, 8, 8	0.014420	0.010480	1.376x
float32	1, 1, 7, 7, 7	0.014200	0.010360	1.371x
float32	2, 3, 10, 10, 10	0.011440	0.010280	1.113x
float32	2, 3, 12, 12, 12	0.012960	0.010440	1.241x
float32	2, 3, 11, 11, 11	0.013960	0.010400	1.342x
float32	1, 1, 18, 18, 18	0.012240	0.007080	1.729x
float32	1, 1, 20, 20, 20	0.016960	0.007480	2.267x
float32	1, 1, 19, 19, 19	0.017040	0.007200	2.367x
float32	2, 4, 10, 18, 34	0.012600	0.007320	1.721x
float32	2, 4, 12, 20, 36	0.016280	0.007640	2.131x
float32	2, 4, 11, 19, 35	0.014240	0.007680	1.854x
bfloat16	1, 1, 6, 6, 6	0.021800	0.010520	2.072x
bfloat16	1, 1, 8, 8, 8	0.038680	0.010840	3.568x
bfloat16	1, 1, 7, 7, 7	0.033840	0.010680	3.169x
bfloat16	2, 3, 10, 10, 10	0.019480	0.010280	1.895x
bfloat16	2, 3, 12, 12, 12	0.022880	0.010280	2.226x
bfloat16	2, 3, 11, 11, 11	0.021960	0.010240	2.145x
bfloat16	1, 1, 18, 18, 18	0.019440	0.006920	2.809x
bfloat16	1, 1, 20, 20, 20	0.025740	0.006880	3.741x
bfloat16	1, 1, 19, 19, 19	0.025120	0.007000	3.589x
bfloat16	2, 4, 10, 18, 34	0.021560	0.007160	3.011x
bfloat16	2, 4, 12, 20, 36	0.027640	0.007320	3.776x
bfloat16	2, 4, 11, 19, 35	0.024240	0.007240	3.348x

Operator	Arithmetic Mean Speedup
reflection_pad3d_backward	2.40x

Files Changed

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

[Dongxu-H]

/test | reflection_pad3d_backward:mthreads

[github-actions]

❌ On-demand test failed

Operator: reflection_pad3d_backward
Runner: mthreads
Backend: mthreads-musa520

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

[Dongxu-H]

/test | reflection_pad3d_backward_moore_threads_specialized:mthreads

[github-actions]

❌ On-demand test failed

Operator: reflection_pad3d_backward_moore_threads_specialized
Runner: mthreads
Backend: mthreads-musa520

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

[Dongxu-H]

/test | reflection_pad3d_backward:mthreads

[github-actions]

⚠️ Cannot proceed with operation

• reviewDecision: skip_reviews
• commitStatus: FAILURE

Reviews are not required for this operation but CI checks must be passing in order to continue

[Dongxu-H]

LGTM

[Dongxu-H]

LGTM