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

[Yukun-Cui]

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

Summary

Add a Moore Threads (MUSA) specialized Triton kernel for adaptive_max_pool3d_backward, overriding the generic implementation via runtime.replace_customized_ops(). This is the backward pass of adaptive_max_pool3d, scattering the gradient of each output window back to its argmax input position.

The kernel uses two paths:

• Divisible case (all spatial axes evenly divide): adaptive windows tile the input disjointly, so each input position belongs to exactly one window. A 3D grid (plane, depth, hw-block) gathers the gradient with a single launch, no zero pass, no atomics, fully coalesced stores.
• Non-divisible case: output windows overlap, so multiple outputs may map to one input. A zero fill pass followed by an atomic scatter accumulates overlapping contributions correctly.

Testing

• Reused the existing upstream accuracy tests tests/test_adaptive_max_pool3d_backward.py (-m adaptive_max_pool3d_backward)
• Validated against reference on the MUSA device; specialization confirmed active via the GEMS_MTHREADS ADAPTIVE_MAX_POOL3D_BACKWARD debug log
• 69 passed, 3 skipped (the 3 skips are output_size=(8,8,8) paired with shape=(1,1,7,7,7), where output_size > input spatial dims — intentional test skips, not failures)
• Non-divisible shapes (7,7,7) and (10,10,10) all PASS via the scatter/atomic path
• Falls back to the generic implementation for unsupported dtype/device/shape (fp64/int64 not supported on Moore Threads hardware)

Performance

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

adaptive_max_pool3d_backward

dtype	Size	Torch Latency (ms)	Gems Latency (ms)	Speedup
float16	1, 1, 4, 4, 4	0.008640	0.004440	1.946x
float16	2, 3, 8, 8, 8	0.006600	0.003000	2.200x
float16	1, 1, 16, 16, 16	0.007240	0.003040	2.382x
float16	2, 8, 4, 4, 4	0.007480	0.003000	2.493x
float32	1, 1, 4, 4, 4	0.008320	0.004600	1.809x
float32	2, 3, 8, 8, 8	0.007560	0.003080	2.455x
float32	1, 1, 16, 16, 16	0.009860	0.003120	3.160x
float32	2, 8, 4, 4, 4	0.007400	0.003000	2.467x
bfloat16	1, 1, 4, 4, 4	0.008600	0.004600	1.870x
bfloat16	2, 3, 8, 8, 8	0.006800	0.003000	2.267x
bfloat16	1, 1, 16, 16, 16	0.007520	0.003040	2.474x
bfloat16	2, 8, 4, 4, 4	0.007620	0.003000	2.540x

Operator	Arithmetic Mean Speedup
adaptive_max_pool3d_backward	2.34x

Files Changed

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

[Dongxu-H]

/test | adaptive_max_pool3d_backward:mthreads

[github-actions]

❌ On-demand test failed

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

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

[Dongxu-H]

/test | adaptive_max_pool3d_backward:mthreads

[Dongxu-H]

/test | adaptive_max_pool3d_backward:mthreads

[Dongxu-H]

/test | adaptive_max_pool3d_backward:mthreads

[Dongxu-H]

/test | adaptive_max_pool3d_backward:mthreads

[Dongxu-H]

LGTM

[Dongxu-H]

LGTM