Migrate DeepGEMM off of pybind11 and on to TORCH_LIBRARY

[cleonard530]

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.

Updates .pyi stub generation to read TORCH_LIBRARY schemas from csrc/ and overlay defaults/types from _C.py.

JitRuntimeHandle in csrc/python_api.cpp preserves the existing _C.get_jit() API through TORCH_LIBRARY. It wraps the existing DeepJIT runtime’s shared pointer in a torch::CustomClassHolder, allowing the dispatcher to return it as a registered custom class. It does not create a new runtime or change kernel execution.

ABI3 note: This migration does not yet enable the stable CPython ABI (abi3) because DeepJIT still uses pybind11 to release the GIL. Since calls through torch.ops already release the GIL, DeepJIT’s GIL handling can be disabled for this integration, removing the remaining pybind11 dependency and allowing abi3 packaging to be restored.

working around the non ABI3 stable CPython was becomming difficult upstream on vllm, so I updated the build to use CPython stable abi. This added a preprocessor macro defined DJ_DISABLE_GIL which is used on my DeepJIT branch deepseek-ai/DeepJIT#11 to not pull in pybind11 and disables the gil release.

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. The _C.py wrapper reconstructs the callback, preserving the existing get_symm_buffer_size_for_mega_moe API and backwards compatibility.

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.
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_moe.hpp

• Added build_symm_buffer_layout, which centralizes the layout calculation and returns the offsets and dimensions needed to construct every symmetric-buffer view. get_symm_buffer_size_for_mega_moe serializes this layout through to_int_list, while the Mega MoE kernel entry points deserialize it and slice the buffer through the shared C++ helper without dispatching through another registered operator. This preserves the original centralized layout calculation, including the num_sms-dependent ring sizing.
• Under pybind, get_symm_buffer_size_for_mega_moe returned an int64_t and a C++ closure. Because TORCH_LIBRARY schemas cannot return closures, the registered operator now returns the byte count and serialized layout as an int[]. _slice_symm_buffer_for_mega_moe(Tensor buffer, int[] layout_info) reconstructs the views. The compatibility wrapper in _C.py captures that layout in a Python closure, preserving the original public API. The slicing operator remains an internal implementation detail and is excluded from _C.__all__.

csrc/python_api.cpp

Replaces the old PYBIND11_MODULE(...) entry point. Each API header declares its schemas with TORCH_LIBRARY_FRAGMENT and registers implementations with TORCH_LIBRARY_IMPL; including those headers in this translation unit instantiates their static registrations when the shared library loads. This file also owns the primary TORCH_LIBRARY(deep_gemm, m) block, which registers the Runtime custom class and get_jit. Finally, REGISTER_EXTENSION(_C_extension) defines the minimal PyInit__C_extension function needed to import the shared library as a Python extension. The Python C API is used only to create that empty module; operators and classes are registered through PyTorch’s library API.

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

  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 test_torch_library.py
  pytest tests/test_coverage_gaps.py  # (local test that fills in the kernel gaps)

All test pass

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

test_mega_gate.py
test_mega_mhc.py
test_mega_moe.py
All six test_sm120_*.py files

and a few test in the other files.

test_layout.py::test_k_grouped_sf_layout_kernels
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::test_sparse_mqa_logits

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

-- extensions --
    [UNSTABLE] [abi3-ok               ] deep_gemm/_C_extension.abi3.so  (stable_shim=0, unstable=97)

This shows that the library is now on the CPython Stable ABI (ABI3)

Performance benchmarks can be found here cleonard530#9. Overall, the performance is on par with the pybind, but take a look at the 3rd case, to see one case where performance could degrade. This is a pathological case that is unlikely to happen in practice, but should be considered.

This PR assumes DeepJIT is compatable with abi3. This is not true on its main branch, so either deepseek-ai/DeepJIT#11 needs to go through first, or minor modifications need to be made here (which would move the build away from abi3 and target specific CPython versions)