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

[Yukun-Cui]

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

Summary

Add a Moore Threads (MUSA) specialized Triton kernel for norm (and its norm.Scalar / norm.ScalarOpt_dim full-reduction dispatches), overriding the generic implementation via the vendor backend replacement system. The full-tensor reduction uses a grid-stride two-pass design (a bounded grid of programs each folding many tiles into a partial, then a small finalize pass) instead of the generic ~sqrt(M) single-load-per-program layout, which collapses occupancy on large tensors. All supported orders (p = 2, +inf, -inf, 0, and general Lp) are handled; partial per-dim reductions defer to the generic path.

Testing

• Reused the existing upstream accuracy tests tests/test_norm.py (-m norm): 45 passed
• Validated on the MUSA device; specialization confirmed active via the GEMS_MTHREADS NORM debug log
• Falls back to the generic implementation for unsupported dtype/device, partial per-dim reductions, and empty inputs (fp64 not supported on Moore Threads hardware)

Performance

Compared against the generic FlagGems implementation on Moore Threads (MUSA). The benchmark exercises the p=2 full-tensor reduction.

norm

dtype	Size	Torch Latency (ms)	Gems Latency (ms)	Speedup
float16	64, 64	0.013000	0.005320	2.444x
float16	1024, 1024	0.014240	0.010480	1.359x
float16	4096, 4096	0.061840	0.061160	1.011x
float16	64, 512, 512	0.061680	0.061200	1.008x
float16	1024, 1024, 1024	1.630240	1.435940	1.135x
float16	10000, 1	0.009200	0.007840	1.173x
float16	10000, 256	0.016880	0.020280	0.832x
float16	10000, 65536	1.020480	0.903000	1.130x
float16	100, 1, 100	0.009260	0.007880	1.175x
float16	100, 256, 100	0.016880	0.020160	0.837x
float16	100, 65536, 100	1.021160	0.902720	1.131x
float32	64, 64	0.007760	0.005760	1.347x
float32	1024, 1024	0.016120	0.013960	1.155x
float32	4096, 4096	0.102420	0.099240	1.032x
float32	64, 512, 512	0.101800	0.098800	1.030x
float32	1024, 1024, 1024	2.994680	2.787960	1.074x
float32	10000, 1	0.010800	0.008120	1.330x
float32	10000, 256	0.025560	0.027720	0.922x
float32	10000, 65536	1.860160	1.729320	1.076x
float32	100, 1, 100	0.008760	0.008080	1.084x
float32	100, 256, 100	0.025520	0.027680	0.922x
float32	100, 65536, 100	1.858160	1.729700	1.074x
bfloat16	64, 64	0.006920	0.005360	1.291x
bfloat16	1024, 1024	0.014120	0.010560	1.337x
bfloat16	4096, 4096	0.061960	0.061360	1.010x
bfloat16	64, 512, 512	0.061840	0.060680	1.019x
bfloat16	1024, 1024, 1024	1.625000	1.437340	1.131x
bfloat16	10000, 1	0.008840	0.007920	1.116x
bfloat16	10000, 256	0.016840	0.020080	0.839x
bfloat16	10000, 65536	1.020700	0.905440	1.127x
bfloat16	100, 1, 100	0.008800	0.007920	1.111x
bfloat16	100, 256, 100	0.016860	0.019960	0.845x
bfloat16	100, 65536, 100	1.019260	0.905440	1.126x

Operator	Arithmetic Mean Speedup
norm	1.13x

Files Changed

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

[Dongxu-H]

/test | norm:mthreads

[Dongxu-H]

LGTM

[Dongxu-H]

/test | norm:mthreads

[github-actions]

❌ On-demand test failed

Operator: norm
Runner: mthreads
Backend: mthreads-musa520

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

[Dongxu-H]

Logger should use MThreads backend-specific format: logging.getLogger(f"flag_gems.runtime.backend._mthreads.ops.{name.split(".")[-1]}") instead of logging.getLogger(name) in backend implementations.

[Dongxu-H]

/test | norm:mthreads

[Dongxu-H]

PR contains large infrastructure changes (removing CODEOWNERS, modifying workflows, adding image-builder and sync-to-kernelgen workflows, modifying backends.yaml, fused_marlin_moe.py, adding unrelated operators and benchmarks) that should not be mixed with operator implementation. Infrastructure changes should be submitted as separate PR. Only operator-specific changes should remain: src/flag_gems/runtime/backend/_mthreads/ops/norm.py.

[Dongxu-H]

/test | zero_:thead

[github-actions]

❌ On-demand test failed

Operator: zero_
Runner: thead
Backend: thead

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

[Dongxu-H]

/test | norm:mthreads

[Dongxu-H]

PR contains large infrastructure changes (removing CODEOWNERS, modifying workflows, adding image-builder and sync-to-kernelgen workflows, modifying backends.yaml, fused_marlin_moe.py, adding unrelated operators and benchmarks) that should not be mixed with operator implementation. Infrastructure changes should be submitted as separate PR. Only operator-specific changes should remain: src/flag_gems/runtime/backend/_mthreads/ops/norm.py.

[Dongxu-H]

/test | norm:mthreads

[Dongxu-H]

PR contains large infrastructure changes (removing CODEOWNERS, modifying workflows, adding image-builder and sync-to-kernelgen workflows, modifying backends.yaml, fused_marlin_moe.py, adding unrelated operators and benchmarks) that should not be mixed with operator implementation. Infrastructure changes should be submitted as separate PR. Only operator-specific changes should remain: src/flag_gems/runtime/backend/_mthreads/ops/norm.py and related init.py update. Merge conflict needs to be resolved.

[Dongxu-H]

/test | norm:mthreads

[Dongxu-H]

/test | norm:mthreads

[Dongxu-H]

LGTM

[Dongxu-H]

/test | norm:mthreads