[Metal] Add reduction and infinity lowering

[GY-Bai]

Problem and change

Metal has no target-specific T.reduce lowering, so normalization and routing reductions used by Qwen dense, Qwen MoE, and DeepSeek-style workloads cannot compile for Apple GPUs. Add a Metal ReduceImpl with fp32 accumulation, uniform barriers, and compile-time checks that keep each XOR butterfly inside a valid simdgroup closure. Also register Metal lowering for tl.infinity; other math intrinsics are unchanged.

_{Files: src/metal/op/reduce.cc, src/metal/op/math.cc, and testing/python/metal/test_metal_reduce.py.}

Before and after

On the same Apple M2, origin/main cannot lower the 16x2 fp32, 32-thread reduction because no Metal implementation is registered. This PR executes reduce_sum and reduce_max on MPS and matches the CPU reference within rtol=1e-5, atol=1e-5.

Performance classification: Performance enablement. This is not an isolated main-versus-PR speedup because main cannot lower the kernel. The following fresh Apple M2 measurements on this PR head use MPS events, three interleaved rounds, 25 warmups, and a 50-launch timing batch; they are downstream evidence enabled by the reduction lowering:

1. Single-token Qwen dense decode: 410.0 us/token (195.0 us/token on the GPU timeline).
2. 64-token Qwen dense decode: 28.3 us/token (23.9 us/token on the GPU timeline), or 14.5x lower synchronized per-token latency.

Validation

cmake -S . -B build
cmake --build build -j8
TILELANG_DISABLE_CACHE=1 python -m pytest testing/python/metal/test_metal_reduce.py -q
31 passed

Twenty-nine compile/lower cases run on every platform; two numerical cases require MPS and skip elsewhere.

Scope

Targets the M1-M4 MSL/simdgroup path; the M5 Metal 4 TensorOps/FP8 path is out of scope. Cross-simdgroup, non-power-of-two or misaligned, batched, and deferred reduction plans remain unsupported.

Appendix: combined campaign results (组合 campaign 结果)

The table below describes the complete Apple M2 optimization campaign, not the isolated contribution or acceptance criteria of this PR. Where origin/main cannot execute the workload correctly, a numerical speedup versus main is undefined; the nearest valid performance baseline is stated explicitly.

Representative workload	origin/main	End-state result	Valid performance comparison
Qwen dense decode	Cannot lower the required Metal reduction	64-token: 28.3 us/token synchronized, 23.9 us/token GPU timeline	14.5x versus the valid single-token path (410.0 us/token)
Fused Qwen/DeepSeek-style MoE block	Required multi-kernel/ABI path is not correct	9 launches, 16.92 ms	1.21x versus the previous valid 17-launch composed path (20.53 ms)
Multi-simdgroup bf16 GEMM	t256 staged output is incorrect	Correct t256 execution	Fragment-C 19.86x; shared-C 14.16x versus correct t32 paths
Software-packed dense fp8 / expert fp4 QMM	Downstream kernels are not present in main	685.7 us / 697.4 us	Dense fp8 16.9x versus the initial campaign kernel; expert fp4 26% faster than its earlier campaign version
End-to-end Mega MoE	Complete path is blocked by missing backend prerequisites	T=64 114.1 ms; T=1 16.7 ms	80.8% / 48.0% lower latency versus the valid full-pipeline baseline

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 1e1d462d-5854-4c7b-accc-f6035dde971b

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[github-actions]

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run pre-commit run --all-files in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀