deepseek-ai/DeepGEMM

Fork 0

mirror of https://github.com/deepseek-ai/DeepGEMM.git synced 2026-11-05 00:47:26 +00:00

[PR #393] Migrate pybind11 to TORCH_LIBRARY #411

New issue

Open

opened 2026-10-26 21:17:24 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:17:24 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/393
Author: @cleonard530
Created: 7/27/2026
Status: 🔄 Open

Base: nv_dev ← Head: migrate_pybind_to_torch_library

📝 Commits (10+)

a3e7b0e migrated off of pybind and onto TORCH_LIBRARY
c91c209 added #if DG..._COMPATABLE guards to the function schema registers to match what was in the legacy code
926d993 updated init.py import to align with legacy code. Removed pybind helper code from generate_pyi.py
c38efbb generate_pyi.py stopped creating the function type hints in the past updates so I updated it to inlcude the type hints
2630e0a updated generate_pyi.py to align more with legacy code
796a48e used _C.py to generate default argumements in gererate_pyi.py since pybind default arguments are no longer there
0749569 reverted back to using TORCH_LIBRARY schema to generate pyi file
fd57d1c 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 main()) in a followup commit, but wanted them here for reference
2a1da8b removed verbose comments and main function from generate_pyi.py
4aff592 added c_py_path for generate_pyi_file in setup.py

📊 Changes

40 files changed (+2255 additions, -1183 deletions)

View changed files

📝 .github/workflows/publish.yml (+3 -6)
📝 csrc/apis/attention.hpp (+125 -34)
📝 csrc/apis/einsum.hpp (+41 -12)
📝 csrc/apis/gemm.hpp (+304 -118)
📝 csrc/apis/hyperconnection.hpp (+28 -5)
📝 csrc/apis/layout.hpp (+98 -26)
📝 csrc/apis/mega.hpp (+400 -96)
📝 csrc/apis/runtime.hpp (+59 -36)
📝 csrc/apis/sm90_mega.hpp (+215 -50)
📝 csrc/jit/device_runtime.hpp (+2 -0)
📝 csrc/jit_kernels/impls/runtime_utils.hpp (+1 -1)
📝 csrc/jit_kernels/impls/sm100_bf16_gemm.hpp (+1 -1)
📝 csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp (+1 -1)
📝 csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp (+1 -1)
📝 csrc/jit_kernels/impls/sm100_fp4_fp4_mega_moe.hpp (+1 -1)
📝 csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp (+1 -1)
📝 csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp (+1 -1)
📝 csrc/jit_kernels/impls/sm100_tf32_hc_prenorm_gemm.hpp (+1 -1)
📝 csrc/jit_kernels/impls/sm120_bf16_gemm.hpp (+1 -1)
📝 csrc/jit_kernels/impls/sm120_bmk_bnk_mn.hpp (+1 -1)

...and 20 more files

📄 Description

Summary

Migrates DeepGEMM's C++ extension from pybind11 (PYBIND11_MODULE + per-header register_apis) to TORCH_LIBRARY / TORCH_LIBRARY_IMPL registration. Ops are now registered as torch.ops.deep_gemm.*.

To preserve the existing Python API (import deep_gemm; deep_gemm._C.fp8_gemm_nt(...)), a new deep_gemm/_C.py shim loads the extension and re-exports the legacy surface. The compiled module is renamed to deep_gemm._C_extension and built as a limited-API (abi3) extension.

Also updates .pyi stub generation to read TORCH_LIBRARY schemas from csrc/ and overlay defaults/types from _C.py, and splits mega MoE symm-buffer sizing/slicing into two separate ops (replacing the old pybind callback return value).

Note: Bumped the abi3 floor to Python 3.10 which is supported by every torch release already in the build matrix (2.4–2.8), and it gets ahead of a floor bump we'll need once DeepGEMM moves to the LibTorch Stable ABI, which requires torch 2.10+ and Python 3.10+.

PR https://github.com/vllm-project/vllm/pull/48962 is used to test this on vLLM.

WARNING:

Before, get_symm_buffer_size_for_mega_moe returned an int and a function (std::tuple<int64_t, std::function...>) but TORCH_LIBRARY doesn't allow returning function objects/closures. Now, it returns an int and the computed buffer layout as an int[] (std::tuple<int64_t, c10::List<int64_t>>), and a separate function, _slice_symm_buffer_for_mega_moe, consumes that int[] instead of the closure. This breaks backwards compatibility for get_symm_buffer_size_for_mega_moe, but this method is not re-exported from deep_gemm/__init__.py (only used internally by SymmBuffer.__init__), and a quick look on GitHub shows no other public repo calls this function directly either.

New files

deep_gemm/_C.py

Python compatibility shim for the old deep_gemm._C module. Responsibilities:

torch.ops.load_library(...) on _C_extension*.so at import time
Re-export torch.ops.deep_gemm.* under the legacy _C names
Restore pybind-era conveniences that are not expressible in TORCH_LIBRARY schemas alone: tuple unpacking (q/kv/(weight, scale) pairs), default arguments, and guarded op binding for #if DG_*_COMPATIBLE build variants
Previously _C was the compiled .so itself; now _C is pure Python and the .so is _C_extension.

csrc/torch_library_utils.hpp

(namespace deep_gemm::torch_utils) Conversion helpers for bridging PyTorch schema types to the C++ types the existing kernel implementations expect — e.g. c10::List<int64_t> → std::tuple<int,int,int> for recipe/head_splits, and optional list → std::vector<int>. Without this, every TORCH_LIBRARY wrapper would duplicate the same list/tuple parsing logic.

csrc/utils/registration.h

Defines the REGISTER_EXTENSION(NAME) macro used by csrc/python_api.cpp. Under TORCH_LIBRARY, ops self-register via static initializers when the .so loads, so there's no real pybind11 module to build — this macro just emits an empty PyInit_ function so the compiled .so still satisfies Python's import machinery (which requires that symbol to exist).

Significantly changed files

csrc/apis/mega.hpp

Added build_symm_buffer_layout, which computes the full SymmBufferLayoutInfo (every sub-buffer's offset/size) once. get_symm_buffer_size_for_mega_moe and fp8_fp4_mega_moe/bf16_mega_moe call it instead of duplicating layout math — the kernel entry points build the layout and slice the buffer directly rather than going through a second registered op. SymmBufferLayoutInfo also gained to_int_list/from_int_list so the layout computed by get_symm_buffer_size_for_mega_moe can be flattened into a plain int[] and handed to _slice_symm_buffer_for_mega_moe, instead of that op recomputing the same layout (including the num_sms-dependent ring sizing) from scratch.
get_symm_buffer_size_for_mega_moe used to return (int64_t, std::function<...>) under pybind; TORCH_LIBRARY can't return closures. It now returns the int plus the computed layout as an int[], and slicing moved to a separate op, _slice_symm_buffer_for_mega_moe(Tensor buffer, int[] layout_info) (leading underscore since it's an internal implementation detail — only deep_gemm/mega/__init__.py's SymmBuffer calls it, it isn't part of the public API).

csrc/python_api.cpp

Replaces the old PYBIND11_MODULE(...) entry point. Ops now self-register via TORCH_LIBRARY_FRAGMENT/STABLE_TORCH_LIBRARY_IMPL static initializers inside each csrc/apis/*.hpp, which run automatically when the .so is loaded — so this file no longer needs to call per-header register_apis(m) functions. It just #includes every apis/*.hpp (to trigger their static registration) and defines an empty PyInit_ symbol via the REGISTER_EXTENSION macro, purely so the compiled .so is still a valid, importable Python extension module — no actual bindings go through the Python C API anymore.

scripts/generate_pyi.py

Previously generated .pyi type hints by regex-parsing the C++ function declarations in csrc/ — but under pybind, the real Python-callable keyword names come from py::arg("...") strings in PYBIND11_MODULE, which pybind doesn't require to match the C++ parameter names. That divergence caused real bugs: e.g. fp8_fp4_paged_mqa_logits's C++ parameter is named fused_kv_cache, but it's bound as py::arg("kv_cache"), so the old generated stub advertised a keyword (fused_kv_cache) that would actually raise TypeError if used — the true kwarg is kv_cache. The script now parses TORCH_LIBRARY schema strings from csrc/ and overlays defaults from an AST parse of deep_gemm/_C.py, so the stub is generated from the real registration contract and entry point instead of an incidental C++ signature.

Test plan/Results

To help verify that each kernel op is still functioning correctly, I created a script, smoke_test_all_ops.py, that performs a minimal call to every op and verifies that any tensor inputs expected to be immutable remain bit-for-bit identical before and after the call.

Checkout commit github.com/deepseek-ai/DeepGEMM@a813b41a1b to see the smoke_test_all_ops.py script that exercises every TORCH_LIBRARY op. I added the results in the comment at the bottom of the commit (they all passed or were skipped because I didn't have the right SM).

I also compared the generated .pyi file before (using C++ signatures directly) and after (using TORCH_LIBRARY schemas) to make sure types are mapped appropriately. The differences can be found here https://github.com/cleonard530/DeepGEMM/pull/5 .

  pytest tests/test_bf16.py
  pytest tests/test_fp8_fp4.py
  pytest tests/test_layout.py
  pytest tests/test_einsum.py
  pytest tests/test_attention.py
  pytest tests/test_hyperconnection.py
  pytest tests/test_legacy.py
  pytest tests/test_coverage_gaps.py  # (local test that fills in the kernel gaps)
  pytest tests/test_lazy_init.py

  python tests/test_mega_moe_sm90.py --num-processes 1

All test pass.

Test that require Blackwell GPUs have not been tested yet. These include

tests/test_mega_moe.py
tests/test_mega_moe_situ.py
tests/test_nvfp4_mega_moe.py

and a few test in

test_layout.py::test_k_grouped_psum_sf_layout_kernels 
test_einsum.py::test_fp8_bhd_hdr_bhr
test_einsum.py::test_fp8_bhd_bhr_hdr 
test_attention.py — selects additional SM100/SM120 formats and varlen cases but still supports Hopper generally

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepGEMM/pull/393 **Author:** [@cleonard530](https://github.com/cleonard530) **Created:** 7/27/2026 **Status:** 🔄 Open **Base:** `nv_dev` ← **Head:** `migrate_pybind_to_torch_library` --- ### 📝 Commits (10+) - [`a3e7b0e`](https://github.com/deepseek-ai/DeepGEMM/commit/a3e7b0ec4118b57ccd98f4825afc50db6b544ead) migrated off of pybind and onto TORCH_LIBRARY - [`c91c209`](https://github.com/deepseek-ai/DeepGEMM/commit/c91c2098a3bba47ae4becf8d02eddd85330d424a) added #if DG..._COMPATABLE guards to the function schema registers to match what was in the legacy code - [`926d993`](https://github.com/deepseek-ai/DeepGEMM/commit/926d9939e09b311e23bd4551bb008c06f74b9226) updated __init__.py import to align with legacy code. Removed pybind helper code from generate_pyi.py - [`c38efbb`](https://github.com/deepseek-ai/DeepGEMM/commit/c38efbb776cf229ac09fbe92344c2eb7b0d0822e) generate_pyi.py stopped creating the function type hints in the past updates so I updated it to inlcude the type hints - [`2630e0a`](https://github.com/deepseek-ai/DeepGEMM/commit/2630e0a5ad2c1cfc236332c1609b7e2244648227) updated generate_pyi.py to align more with legacy code - [`796a48e`](https://github.com/deepseek-ai/DeepGEMM/commit/796a48e0e424b7e64a57bc58d53d1b4a7b583c73) used _C.py to generate default argumements in gererate_pyi.py since pybind default arguments are no longer there - [`0749569`](https://github.com/deepseek-ai/DeepGEMM/commit/074956920d55ffa98368acc8ddc863663f487bb9) reverted back to using TORCH_LIBRARY schema to generate pyi file - [`fd57d1c`](https://github.com/deepseek-ai/DeepGEMM/commit/fd57d1cf4e220ed06c8136f2feb6112732f9cd79) 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 main()) in a followup commit, but wanted them here for reference - [`2a1da8b`](https://github.com/deepseek-ai/DeepGEMM/commit/2a1da8ba79aad454e050800e8d154fbe01a87e64) removed verbose comments and main function from generate_pyi.py - [`4aff592`](https://github.com/deepseek-ai/DeepGEMM/commit/4aff592809d924d382649034e3df10ffc85301d5) added c_py_path for generate_pyi_file in setup.py ### 📊 Changes **40 files changed** (+2255 additions, -1183 deletions) <details> <summary>View changed files</summary> 📝 `.github/workflows/publish.yml` (+3 -6) 📝 `csrc/apis/attention.hpp` (+125 -34) 📝 `csrc/apis/einsum.hpp` (+41 -12) 📝 `csrc/apis/gemm.hpp` (+304 -118) 📝 `csrc/apis/hyperconnection.hpp` (+28 -5) 📝 `csrc/apis/layout.hpp` (+98 -26) 📝 `csrc/apis/mega.hpp` (+400 -96) 📝 `csrc/apis/runtime.hpp` (+59 -36) 📝 `csrc/apis/sm90_mega.hpp` (+215 -50) 📝 `csrc/jit/device_runtime.hpp` (+2 -0) 📝 `csrc/jit_kernels/impls/runtime_utils.hpp` (+1 -1) 📝 `csrc/jit_kernels/impls/sm100_bf16_gemm.hpp` (+1 -1) 📝 `csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp` (+1 -1) 📝 `csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp` (+1 -1) 📝 `csrc/jit_kernels/impls/sm100_fp4_fp4_mega_moe.hpp` (+1 -1) 📝 `csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp` (+1 -1) 📝 `csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp` (+1 -1) 📝 `csrc/jit_kernels/impls/sm100_tf32_hc_prenorm_gemm.hpp` (+1 -1) 📝 `csrc/jit_kernels/impls/sm120_bf16_gemm.hpp` (+1 -1) 📝 `csrc/jit_kernels/impls/sm120_bmk_bnk_mn.hpp` (+1 -1) _...and 20 more files_ </details> ### 📄 Description ## Summary Migrates DeepGEMM's C++ extension from pybind11 (`PYBIND11_MODULE` + per-header `register_apis`) to `TORCH_LIBRARY` / `TORCH_LIBRARY_IMPL` registration. Ops are now registered as `torch.ops.deep_gemm.*`. To preserve the existing Python API (`import deep_gemm; deep_gemm._C.fp8_gemm_nt(...)`), a new `deep_gemm/_C.py` shim loads the extension and re-exports the legacy surface. The compiled module is renamed to `deep_gemm._C_extension` and built as a limited-API (abi3) extension. Also updates `.pyi` stub generation to read `TORCH_LIBRARY` schemas from `csrc/` and overlay defaults/types from `_C.py`, and splits mega MoE symm-buffer sizing/slicing into two separate ops (replacing the old pybind callback return value). Note: Bumped the abi3 floor to Python 3.10 which is supported by every torch release already in the build matrix (2.4–2.8), and it gets ahead of a floor bump we'll need once DeepGEMM moves to the LibTorch Stable ABI, which requires torch 2.10+ and Python 3.10+. PR https://github.com/vllm-project/vllm/pull/48962 is used to test this on vLLM. ## WARNING: Before, `get_symm_buffer_size_for_mega_moe` returned an `int` and a `function` (`std::tuple<int64_t, std::function...>`) but `TORCH_LIBRARY` doesn't allow returning function objects/closures. Now, it returns an `int` and the computed buffer layout as an `int[]` (`std::tuple<int64_t, c10::List<int64_t>>`), and a separate function, `_slice_symm_buffer_for_mega_moe`, consumes that `int[]` instead of the closure. **This breaks backwards compatibility** for `get_symm_buffer_size_for_mega_moe`, but this method is not re-exported from `deep_gemm/__init__.py` (only used internally by `SymmBuffer.__init__`), and a quick look on GitHub shows no other public repo calls this function directly either. ## New files **deep_gemm/_C.py** Python compatibility shim for the old `deep_gemm._C` module. Responsibilities: `torch.ops.load_library(...)` on `_C_extension*.so` at import time Re-export `torch.ops.deep_gemm.*` under the legacy `_C` names Restore pybind-era conveniences that are not expressible in `TORCH_LIBRARY` schemas alone: tuple unpacking (`q`/`kv`/`(weight, scale)` pairs), default arguments, and guarded op binding for `#if DG_*_COMPATIBLE` build variants Previously `_C` was the compiled `.so` itself; now `_C` is pure Python and the `.so` is `_C_extension`. **csrc/torch_library_utils.hpp** (namespace `deep_gemm::torch_utils`) Conversion helpers for bridging PyTorch schema types to the C++ types the existing kernel implementations expect — e.g. `c10::List<int64_t>` → `std::tuple<int,int,int>` for `recipe`/`head_splits`, and optional list → `std::vector<int>`. Without this, every `TORCH_LIBRARY` wrapper would duplicate the same list/tuple parsing logic. **csrc/utils/registration.h** Defines the REGISTER_EXTENSION(NAME) macro used by csrc/python_api.cpp. Under TORCH_LIBRARY, ops self-register via static initializers when the .so loads, so there's no real pybind11 module to build — this macro just emits an empty PyInit_<NAME> function so the compiled .so still satisfies Python's import machinery (which requires that symbol to exist). ## Significantly changed files **csrc/apis/mega.hpp** - Added `build_symm_buffer_layout`, which computes the full `SymmBufferLayoutInfo` (every sub-buffer's offset/size) once. `get_symm_buffer_size_for_mega_moe` and `fp8_fp4_mega_moe`/`bf16_mega_moe` call it instead of duplicating layout math — the kernel entry points build the layout and slice the buffer directly rather than going through a second registered op. `SymmBufferLayoutInfo` also gained `to_int_list`/`from_int_list` so the layout computed by `get_symm_buffer_size_for_mega_moe` can be flattened into a plain `int[]` and handed to `_slice_symm_buffer_for_mega_moe`, instead of that op recomputing the same layout (including the `num_sms`-dependent ring sizing) from scratch. - `get_symm_buffer_size_for_mega_moe` used to return (`int64_t`, `std::function<...>`) under pybind; `TORCH_LIBRARY` can't return closures. It now returns the int plus the computed layout as an `int[]`, and slicing moved to a separate op, `_slice_symm_buffer_for_mega_moe(Tensor buffer, int[] layout_info)` (leading underscore since it's an internal implementation detail — only `deep_gemm/mega/__init__.py`'s `SymmBuffer` calls it, it isn't part of the public API). **csrc/python_api.cpp** Replaces the old `PYBIND11_MODULE(...)` entry point. Ops now self-register via `TORCH_LIBRARY_FRAGMENT`/`STABLE_TORCH_LIBRARY_IMPL` static initializers inside each `csrc/apis/*.hpp`, which run automatically when the `.so` is loaded — so this file no longer needs to call per-header `register_apis(m)` functions. It just `#includes` every `apis/*.hpp` (to trigger their static registration) and defines an empty `PyInit_` symbol via the `REGISTER_EXTENSION` macro, purely so the compiled `.so` is still a valid, importable Python extension module — no actual bindings go through the Python C API anymore. **scripts/generate_pyi.py** Previously generated `.pyi` type hints by regex-parsing the C++ function declarations in `csrc/` — but under pybind, the real Python-callable keyword names come from `py::arg("...")` strings in `PYBIND11_MODULE`, which pybind doesn't require to match the C++ parameter names. That divergence caused real bugs: e.g. `fp8_fp4_paged_mqa_logits`'s C++ parameter is named `fused_kv_cache`, but it's bound as `py::arg("kv_cache")`, so the old generated stub advertised a keyword (`fused_kv_cache`) that would actually raise `TypeError` if used — the true kwarg is `kv_cache`. The script now parses TORCH_LIBRARY schema strings from `csrc/` and overlays defaults from an AST parse of `deep_gemm/_C.py`, so the stub is generated from the real registration contract and entry point instead of an incidental C++ signature. ## Test plan/Results 1. To help verify that each kernel op is still functioning correctly, I created a script, `smoke_test_all_ops.py`, that performs a minimal call to every op and verifies that any tensor inputs expected to be immutable remain bit-for-bit identical before and after the call. Checkout commit https://github.com/deepseek-ai/DeepGEMM/commit/a813b41a1b8b9a1ef96dc37c5b35a3cc4d7a0150 to see the `smoke_test_all_ops.py` script that exercises every TORCH_LIBRARY op. I added the results in the comment at the bottom of the commit (they all passed or were skipped because I didn't have the right SM). 2. I also compared the generated .pyi file before (using C++ signatures directly) and after (using TORCH_LIBRARY schemas) to make sure types are mapped appropriately. The differences can be found here https://github.com/cleonard530/DeepGEMM/pull/5 . 3. ``` pytest tests/test_bf16.py pytest tests/test_fp8_fp4.py pytest tests/test_layout.py pytest tests/test_einsum.py pytest tests/test_attention.py pytest tests/test_hyperconnection.py pytest tests/test_legacy.py pytest tests/test_coverage_gaps.py # (local test that fills in the kernel gaps) pytest tests/test_lazy_init.py python tests/test_mega_moe_sm90.py --num-processes 1 ``` All test pass. Test that require Blackwell GPUs have not been tested yet. These include ``` tests/test_mega_moe.py tests/test_mega_moe_situ.py tests/test_nvfp4_mega_moe.py ``` and a few test in ``` test_layout.py::test_k_grouped_psum_sf_layout_kernels test_einsum.py::test_fp8_bhd_hdr_bhr test_einsum.py::test_fp8_bhd_bhr_hdr test_attention.py — selects additional SM100/SM120 formats and varlen cases but still supports Hopper generally ``` --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-26 21:17:24 +00:00

bsakaguchi referenced this issue

2026-10-26 21:17:39 +00:00

[PR #411] feat: add SM90 FP8 x MXFP4 MegaMoE support #427

bsakaguchi referenced this issue

2026-10-26 21:18:00 +00:00

[PR #440] Add a fused SM90 FP8 MegaMoE: one kernel, a lag-ordered schedule and a ring activation pool #449

bsakaguchi referenced this issue

2026-10-26 21:18:18 +00:00