Migrate DeepGEMM from pybind11 to TORCH_LIBRARY op registration

[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 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).

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 just returns an int, and a separate function, _slice_symm_buffer_for_mega_moe, is used 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.

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, slice_symm_buffer_for_mega_moe, and fp8_fp4_mega_moe/bf16_mega_moe all 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.
• 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 just the int, and slicing moved to a separate op, _slice_symm_buffer_for_mega_moe (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 4cb88be 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 Gen pyi test #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!

[cleonard530]

sfa and sfb are here because TORCH_LIBRARY schemas can’t express pybind-style (tensor, scale) pairs. They need flat tensor arguments.

This shows up a lot in this migration.

[cleonard530]

TORCH_LIBRARY cannot return a callback like pybind11 can, so a lamda function on the python side is created.

[cleonard530]

Wrapper is needed since TORCH_LIBRARY int? maps to optional<int64_t>, not optional<int>

[cleonard530]

Allows us to use torch::... instead of at::...

[cleonard530]

Look into another header (e.g torch/all.h) instead of this wrapper

[cleonard530]

Defines a shorthand for binding adapter functions in deep_gemm::torch_registration

[cleonard530]

Helper functions to turn c10::List<int64_t> (TORCH_LIBRARY allowed type) to std::tuples<int,...> (existing kernel APIs) are defined in this helper file.

[cleonard530]

unused import

[cleonard530]

This file loads the compiled _C_extension*.so via torch.ops.load_library, wraps torch.ops.deep_gemm with Python shims that preserve the old API (tuple unpacking, guarded exports), and re-exports those ops as the deep_gemm._C module.

This keeps import _C consistency with legacy code

[cleonard530]

This file used to build the .pyi file by parsing the C++ functions. Now it used the TORCH_LIBRARY schema and _C.py (for default values) to build the .pyi file with type hints.

[cleonard530]

The differemce in the generated file can be seen here: 02ed6bc

[cleonard530]

never used

[cleonard530]

never used

[cleonard530]

Doesn't do anything in this while loop

[cleonard530]

Get default arguments from _C.py

[cleonard530]

fp8_fp4_mega_moe and bf16_mega_moe now use the internal build_symm_buffer_layout to get the num_required_bytes and slice (through slice_symm_buffer_from_layout) instead of the registered torch ops get_symm_buffer_size_for_mega_moe, which used to return a callback but is unable to with TORCH_LIBRARY (it was allowed on pybind). This op is still registered and used in deep_gemm/mega/__init__.py as well as slice_symm_buffer_for_mega_moe (to replace the callback) which is a small wrapper around slice_symm_buffer_from_layout.

[cleonard530]

Need to remove this and use dtype

[cleonard530]

make shorter (e.g torch_utils)