Migrate pybind11 to TORCH_LIBRARY by cleonard530 · Pull Request #393 · deepseek-ai/DeepGEMM

cleonard530 · 2026-07-27T12:31:33Z

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 vllm-project/vllm#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 a813b41 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 Gen pyi test cleonard530/DeepGEMM#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

All test pass!