Repository navigation
Replace the tvm-ffi binding with TORCH_LIBRARY ops - #101
Merged
Merged
Conversation
Register the `_C.so` entry points as `torch.ops.deep_gemm.*` instead of tvm-ffi functions. Tensors cross the boundary as `at::Tensor`, so the DLPack conversions in both directions go away. The library has no CPython dependency (no pybind11, no Py* symbols), so a single wheel still serves every Python version. Bump DeepJIT to 3732a3b and build with `DJ_DISABLE_GIL`, since the dispatcher already releases the GIL. Drop the `apache-tvm-ffi` dependency. Closure-returning symmetric-buffer APIs are split into a size op and a slice op; the Python API is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Fridge003
approved these changes
Oct 13, 2026
zcnrex
added a commit
to zcnrex/DeepGEMM
that referenced
this pull request
Oct 14, 2026
dev replaced the TVM-FFI binding with TORCH_LIBRARY ops (sgl-project#101). Port the sync's binding work onto it: - Epilogue classes cross as (kind, alpha, sfd) into every GEMM/einsum op; epilogue_sfd is a mutable schema argument. - Attention ops follow upstream's API; skip-head-mid, the legacy MQA logits ops and get_bf16_mega_gate_config are not registered. - shutdown and the locality-domain ops are registered; dtypes cross as ScalarType, so FP8/FP4 allocations no longer need a byte view. - MLOPart reaches Python through ctypes callbacks installed by deep_gemm/locality_domain.py, keeping _C free of CPython. - deep_jit stays at dev's 3732a3b, a descendant of upstream's 2efdab4. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1 of 2 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Register the
_C.soentry points astorch.ops.deep_gemm.*(csrc/torch_library_api.cpp) instead of tvm-ffi functions. Same op set and argument order, explicit schemas with mutation annotations, and tensors cross asat::Tensor, so the DLPack round trips are gone.Py*symbols, no pybind11, nolibtvm_ffi.so), so one wheel still runs on every Python version. It is still built against one PyTorch version.DJ_DISABLE_GIL(the dispatcher already releases the GIL), which drops its pybind11 dependency.get_symm_buffer_size_for_{,sm90_}mega_moereturned closures, which schemas cannot express; they are split into a size op and aslice_symm_buffer_for_*op. The Python API is unchanged.apache-tvm-ffiis no longer a dependency;sgl_deep_gemm/build_ext.pybuilds the single translation unit with a plain C++ compiler.Validated on 1x B200 (torch 2.13, CUDA 13.0):
run_tests.sh --release --skip-sanitizer: all single-GPU tests pass, including mega_gate, mega_mhc and the single-rank mega_moe tests.test_mega_moe.py/test_mega_moe_sgl.pyfail identically on the tvm-ffi baseline (DeepEP needs more than one rank); multi-rank and SM90 tests were not run.🤖 Generated with Claude Code