[10b/n] Migrate custom all-reduce, DeepSeek V4 fused MLA, MiniMax reduce-RMS, and MXFP8 MoE to libtorch stable ABI

[cleonard530]

Purpose

Continues the libtorch stable ABI migration by moving several kernels out of legacy _C and into _C_stable_libtorch.

This PR migrates custom all-reduce, DeepSeek V4 fused MLA, MiniMax reduce-RMS, and MXFP8 MoE kernels from legacy _C to _C_stable_libtorch, converting host code to stable torch APIs and moving CMake/bindings accordingly. QuickReduce stays on legacy _C for ROCm-only builds

cc @janeyx99 @Harry-Chen

See 7fc426c to see diffs. A few files look like created/deleted instead of moved otherwise.

Test Plan

pytest tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py 
pytest tests/distributed/test_custom_all_reduce.py 
pytest tests/kernels/core/test_minimax_reduce_rms.py 
pytest tests/kernels/moe/test_cutlass_mxfp8_grouped_mm.py (needs Blackwell GPU to run)

Test Result

(Same test is failing on main branch) uv ---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

---

Migration progress using the Audit Python extension torch-abi-audit:

main branch

  -- extensions --
    [UNSTABLE] [abi3-ok               ] _C.abi3.so  (stable_shim=0, unstable=78)
    [STABLE  ] [abi3-ok               ] _C_stable_libtorch.abi3.so  (stable_shim=81, unstable=0)
    [UNSTABLE] [abi3-ok               ] _flashmla_C.abi3.so  (stable_shim=0, unstable=72)
    [UNSTABLE] [abi3-ok               ] _flashmla_extension_C.abi3.so  (stable_shim=0, unstable=68)
    [UNSTABLE] [abi3-ok               ] _moe_C.abi3.so  (stable_shim=0, unstable=83)
    [NO-TORCH] [abi3-ok               ] cumem_allocator.abi3.so
    [NO-TORCH] [abi3-ok               ] spinloop.abi3.so
    [UNSTABLE] [uses-private-api      ] third_party/deep_gemm/_C.cpython-312-x86_64-linux-gnu.so  (stable_shim=0, unstable=59)
    [UNSTABLE] [abi3-ok               ] vllm_flash_attn/_vllm_fa2_C.abi3.so  (stable_shim=0, unstable=85)
    [UNSTABLE] [abi3-ok               ] vllm_flash_attn/_vllm_fa3_C.abi3.so  (stable_shim=0, unstable=81)

This branch

-- extensions --
    [UNSTABLE] [abi3-ok               ] _C.abi3.so  (stable_shim=0, unstable=78)
    [STABLE  ] [abi3-ok               ] _C_stable_libtorch.abi3.so  (stable_shim=81, unstable=0)
    [UNSTABLE] [abi3-ok               ] _flashmla_C.abi3.so  (stable_shim=0, unstable=72)
    [UNSTABLE] [abi3-ok               ] _flashmla_extension_C.abi3.so  (stable_shim=0, unstable=68)
    [UNSTABLE] [abi3-ok               ] _moe_C.abi3.so  (stable_shim=0, unstable=83)
    [NO-TORCH] [abi3-ok               ] cumem_allocator.abi3.so
    [NO-TORCH] [abi3-ok               ] spinloop.abi3.so
    [UNSTABLE] [uses-private-api      ] third_party/deep_gemm/_C.cpython-312-x86_64-linux-gnu.so  (stable_shim=0, unstable=58)
    [UNSTABLE] [uses-private-api      ] third_party/deep_gemm/_C.cpython-313-x86_64-linux-gnu.so  (stable_shim=0, unstable=59)
    [UNSTABLE] [abi3-ok               ] vllm_flash_attn/_vllm_fa2_C.abi3.so  (stable_shim=0, unstable=85)
    [UNSTABLE] [abi3-ok               ] vllm_flash_attn/_vllm_fa3_C.abi3.so  (stable_shim=0, unstable=81)

No symbols removed

[cleonard530]

Had to redefine _is_weak_contiguous() because the stable wrapper does not expose storage()

[janeyx99]

how come no move for this one

[cleonard530]

This is what I mentioned in the comments. I guess because there were too many updates the diff between this and main was too large so GitHub treated it as a deleted/created file instead of a move. If you checkout the commit linked in the description though you can see the diff better

[janeyx99]

same here, why no move?

[cleonard530]

Same as above