Migrate DeepGEMM from pybind11 to TORCH_LIBRARY op registration - #2

Draft

cleonard530 wants to merge 29 commits into

nv_devfrom

migrate_pybind_to_torch_library

Commits on Sep 8, 2026

Add SM90 FP8 MegaMoE support (#383)

AichenF
and
liz-badada
authored
572557e

Commits on Sep 9, 2026

migrated off of pybind and onto TORCH_LIBRARY
cleonard530
committed
a3e7b0e
added #if DG..._COMPATABLE guards to the function schema registers to match what was in the legacy code
cleonard530
committed
c91c209
updated __init__.py import to align with legacy code. Removed pybind helper code from generate_pyi.py
cleonard530
committed
926d993
generate_pyi.py stopped creating the function type hints in the past updates so I updated it to inlcude the type hints
cleonard530
committed
c38efbb
updated generate_pyi.py to align more with legacy code
cleonard530
committed
2630e0a
used _C.py to generate default argumements in gererate_pyi.py since pybind default arguments are no longer there
cleonard530
committed
796a48e
reverted back to using TORCH_LIBRARY schema to generate pyi file
cleonard530
committed
0749569
made some updates to the generate_pyi.py file to remove aliases and other minor formatting issues. Also added comments with examples to help describe each step. These will be removed (as well as ma…
cleonard530
committed
fd57d1c
removed verbose comments and main function from generate_pyi.py
cleonard530
committed
2a1da8b
added c_py_path for generate_pyi_file in setup.py
cleonard530
committed
4aff592
Refactor mega_moe to build symm buffer layout once per kernel launch
cleonard530
committed
ee76243
Match the FA/TORCH_LIBRARY packaging pattern: ops still register via TORCH_LIBRARY; deep_gemm/_C.py can keep using torch.ops.load_library.
cleonard530
committed
75d4a62
rebased onto nv_dev, removed duplicate slice_symm_buffer_for_mega_moe function that was leftover from migration, replaces torch_compat with torch/all.h, and updated the sm120 files to use torch/all…
cleonard530
committed
927c5a7
Rename TORCH_LIBRARY schemas to match C++ parameter names; drop torch_library_macros
cleonard530
committed
a71e0a8
Fix test_paged_mqa_logits kwarg to match fused_kv_cache rename
cleonard530
committed
0022d45
reverting pytorch apis back to the original names instead of trying to align it with the C++ names
cleonard530
committed
94d3bfb
Mark cumulative_local_expert_recv_stats as mutable in mega MoE schemas.
cleonard530
committed
82c7c96
Align the abi3 floor with the release matrix to stop colliding wheel filenames on release uploads. Bump the floor to cp310, which every torch version in the matrix (2.4-2.8) already supports.
cleonard530
committed
944a443
Make required in fp8_einsum and k_grouped_fp8_gemm_*_contiguous schemas. These were declared optional but the C++ impl always called unconditionally, so would have crashed anyway; the Python wra…
cleonard530
committed
d689a67
Avoid recomputing the mega MoE symm-buffer layout twice per allocation by having get_symm_buffer_size_for_mega_moe return the computed layout as an int[] alongside num_bytes, which _slice_symm_buff…
cleonard530
committed
5099788
Use int[N] instead of int[] in TORCH_LIBRARY schemas for fixed-arity params like recipe/head_splits, so the schema itself is the source of truth for stub generation instead of a hand-maintained per…
cleonard530
committed
5f39e32
Generalize _merge_named_pairs to cover asymmetric-optional tensor/scale-factor pairs like q/q_sf, replacing the special-cased _apply_q_qsf_merge, and sort the csrc/ file scan so op ordering in the …
cleonard530
committed
5c74217
replaced c10::List that I added with std::vector. It doesn't really matter for this migration but it makes the migration to torch abi stable easier
cleonard530
committed
a9a9576
Make _bind_guarded_ops assert on partial guard-group matches instead of silently dropping the whole group, so a future C++ #if drift between grouped ops fails loudly at import with the exact op nam…
cleonard530
committed
3ea26ef
Preserve NVFP4 MegaMoE support after TORCH_LIBRARY refactor
cleonard530
committed
e6e1dcf
fixed comments and changed torch/python.h to torch/all.h in new csrc/jit_kernels/impls/sm100_fp4_fp4_mega_moe.hpp file
cleonard530
committed
9b82e2e
Migrate SM90 MegaMoE APIs to TORCH_LIBRARY
cleonard530
committed
343826c
removed make_tuple and replaced it with curly brackets {}. The return type is explicit so the make_tuple was unnessesary, and this keeps it consistent. with sm90_mega
cleonard530
committed
043dbda

Navigation Menu

Migrate DeepGEMM from pybind11 to TORCH_LIBRARY op registration - #2

Migrate DeepGEMM from pybind11 to TORCH_LIBRARY op registration#2
cleonard530 wants to merge 29 commits into
nv_devfrom
migrate_pybind_to_torch_library

Commits on Sep 8, 2026

Add SM90 FP8 MegaMoE support (#383)

Commits on Sep 9, 2026

migrated off of pybind and onto TORCH_LIBRARY

added #if DG..._COMPATABLE guards to the function schema registers to match what was in the legacy code

updated init.py import to align with legacy code. Removed pybind helper code from generate_pyi.py

generate_pyi.py stopped creating the function type hints in the past updates so I updated it to inlcude the type hints

updated generate_pyi.py to align more with legacy code

used _C.py to generate default argumements in gererate_pyi.py since pybind default arguments are no longer there

reverted back to using TORCH_LIBRARY schema to generate pyi file

made some updates to the generate_pyi.py file to remove aliases and other minor formatting issues. Also added comments with examples to help describe each step. These will be removed (as well as ma…

removed verbose comments and main function from generate_pyi.py

added c_py_path for generate_pyi_file in setup.py

Refactor mega_moe to build symm buffer layout once per kernel launch

Match the FA/TORCH_LIBRARY packaging pattern: ops still register via TORCH_LIBRARY; deep_gemm/_C.py can keep using torch.ops.load_library.

rebased onto nv_dev, removed duplicate slice_symm_buffer_for_mega_moe function that was leftover from migration, replaces torch_compat with torch/all.h, and updated the sm120 files to use torch/all…

Rename TORCH_LIBRARY schemas to match C++ parameter names; drop torch_library_macros

Fix test_paged_mqa_logits kwarg to match fused_kv_cache rename

reverting pytorch apis back to the original names instead of trying to align it with the C++ names

Mark cumulative_local_expert_recv_stats as mutable in mega MoE schemas.

Align the abi3 floor with the release matrix to stop colliding wheel filenames on release uploads. Bump the floor to cp310, which every torch version in the matrix (2.4-2.8) already supports.

Make required in fp8_einsum and k_grouped_fp8_gemm_*_contiguous schemas. These were declared optional but the C++ impl always called unconditionally, so would have crashed anyway; the Python wra…

Avoid recomputing the mega MoE symm-buffer layout twice per allocation by having get_symm_buffer_size_for_mega_moe return the computed layout as an int[] alongside num_bytes, which _slice_symm_buff…

Use int[N] instead of int[] in TORCH_LIBRARY schemas for fixed-arity params like recipe/head_splits, so the schema itself is the source of truth for stub generation instead of a hand-maintained per…

Generalize _merge_named_pairs to cover asymmetric-optional tensor/scale-factor pairs like q/q_sf, replacing the special-cased _apply_q_qsf_merge, and sort the csrc/ file scan so op ordering in the …

replaced c10::List that I added with std::vector. It doesn't really matter for this migration but it makes the migration to torch abi stable easier

Make _bind_guarded_ops assert on partial guard-group matches instead of silently dropping the whole group, so a future C++ #if drift between grouped ops fails loudly at import with the exact op nam…

Preserve NVFP4 MegaMoE support after TORCH_LIBRARY refactor

fixed comments and changed torch/python.h to torch/all.h in new csrc/jit_kernels/impls/sm100_fp4_fp4_mega_moe.hpp file

Migrate SM90 MegaMoE APIs to TORCH_LIBRARY

removed make_tuple and replaced it with curly brackets {}. The return type is explicit so the make_tuple was unnessesary, and this keeps it consistent. with sm90_mega