From 4b4650b426d663e8bac4774989b5132a58eaf85c Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Mon, 12 Oct 2026 14:55:43 +0000 Subject: [PATCH 1/3] Support MLOPart imports from vendored package names Co-authored-by: codex Signed-off-by: Chris Leonard --- csrc/apis/config.hpp | 4 +++- csrc/runtime/mlopart.hpp | 6 ++++-- csrc/runtime/python_package.hpp | 17 +++++++++++++++++ deep_gemm/_C.py | 6 +++++- 4 files changed, 29 insertions(+), 4 deletions(-) create mode 100644 csrc/runtime/python_package.hpp diff --git a/csrc/apis/config.hpp b/csrc/apis/config.hpp index a14b91cf20..8d17228966 100644 --- a/csrc/apis/config.hpp +++ b/csrc/apis/config.hpp @@ -3,12 +3,14 @@ #include #include "../jit_kernels/heuristics/runtime.hpp" +#include "../runtime/python_package.hpp" #include "../runtime/runtime.hpp" TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { using namespace deep_gemm; - m.def("init(str library_root_path) -> ()", [&](const std::string& library_root_path) { + m.def("init(str library_root_path, str python_package_name='deep_gemm') -> ()", [&](const std::string& library_root_path, const std::string& python_package_name) { init_jit(library_root_path); + python_package::set_package_name(python_package_name); }); m.def("shutdown() -> ()", []() { runtime = deep_jit::LazyInit(nullptr); diff --git a/csrc/runtime/mlopart.hpp b/csrc/runtime/mlopart.hpp index 2d30ed2540..eed54ff6a9 100644 --- a/csrc/runtime/mlopart.hpp +++ b/csrc/runtime/mlopart.hpp @@ -11,6 +11,7 @@ #include "../utils/exception.hpp" #include "jit.hpp" +#include "python_package.hpp" namespace deep_gemm::mlopart { @@ -69,7 +70,7 @@ static PyObjectRef checked(PyObject* ptr) { return PyObjectRef(ptr); } -// `deep_gemm.utils.mlopart.(*args)`, stealing the references of `args` +// `.utils.mlopart.(*args)`, stealing the references of `args` // NOTES: the caller must hold the GIL static PyObjectRef call(const char* name, const std::initializer_list& args) { PyObjectRef tuple(PyTuple_New(static_cast(args.size()))); @@ -85,7 +86,8 @@ static PyObjectRef call(const char* name, const std::initializer_list check_python_error(); DG_HOST_ASSERT(tuple.get() != nullptr); - const auto module = checked(PyImport_ImportModule("deep_gemm.utils.mlopart")); + const auto module_name = python_package::module_name("utils.mlopart"); + const auto module = checked(PyImport_ImportModule(module_name.c_str())); const auto func = checked(PyObject_GetAttrString(module.get(), name)); return checked(PyObject_Call(func.get(), tuple.get(), nullptr)); } diff --git a/csrc/runtime/python_package.hpp b/csrc/runtime/python_package.hpp new file mode 100644 index 0000000000..4b31700978 --- /dev/null +++ b/csrc/runtime/python_package.hpp @@ -0,0 +1,17 @@ +#pragma once + +#include + +namespace deep_gemm::python_package { + +inline std::string package_name = "deep_gemm"; + +inline void set_package_name(const std::string& name) { + package_name = name; +} + +inline std::string module_name(const char* relative_name) { + return package_name + "." + relative_name; +} + +} // namespace deep_gemm::python_package diff --git a/deep_gemm/_C.py b/deep_gemm/_C.py index 05677ad635..bb909bb785 100644 --- a/deep_gemm/_C.py +++ b/deep_gemm/_C.py @@ -36,7 +36,7 @@ def _bind_guarded_ops(*names): globals().update({name: getattr(_torch_ops, name) for name in names}) -init = _torch_ops.init +_cpp_init = _torch_ops.init shutdown = _torch_ops.shutdown set_num_sms = _torch_ops.set_num_sms get_num_sms = _torch_ops.get_num_sms @@ -53,6 +53,10 @@ def _bind_guarded_ops(*names): cublaslt_gemm_tt = _torch_ops.cublaslt_gemm_tt +def init(library_root_path): + _cpp_init(library_root_path, __package__) + + def set_block_size_multiple_of(value): if isinstance(value, int): return _torch_ops.set_block_size_multiple_of([value]) From a8b530897ee925ca50bdf1d6872822e127786ef6 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Mon, 12 Oct 2026 17:26:15 +0000 Subject: [PATCH 2/3] better error handeling for mlopart if SM100 is not available. Signed-off-by: Chris Leonard --- deep_gemm/utils/mlopart.py | 12 ++++++++++-- tests/test_locality_domain.py | 6 ++++-- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/deep_gemm/utils/mlopart.py b/deep_gemm/utils/mlopart.py index 3f5b923606..af516c4f5b 100644 --- a/deep_gemm/utils/mlopart.py +++ b/deep_gemm/utils/mlopart.py @@ -19,6 +19,10 @@ PR_SET_PDEATHSIG = 1 # `` +class MLOPartUnavailable(RuntimeError): + pass + + class MLOPart: """ An MLOPart client that allocates localized memory. @@ -96,7 +100,8 @@ def serve(sock: socket.socket) -> None: check(cuda.cuInit(0)) num_devices = ctypes.c_int() check(cuda.cuDeviceGetCount(ctypes.byref(num_devices))) - assert num_devices.value == NUM_LOCALITY_DOMAINS, f'{num_devices.value} MLOPart devices' + if num_devices.value != NUM_LOCALITY_DOMAINS: + raise MLOPartUnavailable(f'{num_devices.value} MLOPart devices') sock.sendall(READY) while request := sock.recv(REQUEST.size, socket.MSG_WAITALL): @@ -131,4 +136,7 @@ def main(sock_fd: int, device: str) -> None: if __name__ == '__main__': - main(int(sys.argv[1]), sys.argv[2]) + try: + main(int(sys.argv[1]), sys.argv[2]) + except MLOPartUnavailable: + pass diff --git a/tests/test_locality_domain.py b/tests/test_locality_domain.py index a025e1cef6..5a3aaf006d 100644 --- a/tests/test_locality_domain.py +++ b/tests/test_locality_domain.py @@ -135,9 +135,10 @@ def test_without_mlopart() -> None: # MPS predating MLOPart ignores `-mlopart`: emulate it by stripping the flag before the real control CLI # NOTES: through the torch dispatcher, the Python `AssertionError` from MLOPart surfaces as a `RuntimeError` real = shutil.which('nvidia-cuda-mps-control') - script = ('import warnings, torch, deep_gemm\n' + script = ('import warnings, torch\n' 'with warnings.catch_warnings(record=True) as caught:\n' ' warnings.simplefilter("always")\n' + ' import deep_gemm\n' ' assert not deep_gemm.is_localization_available()\n' 'assert len(caught) == 1 and "MLOPart is unavailable" in str(caught[0].message)\n' 'assert deep_gemm.locality_domain.get_sm_locality_domains().numel() == torch.cuda.get_device_properties(0).multi_processor_count\n' @@ -153,7 +154,8 @@ def test_without_mlopart() -> None: fake.chmod(0o755) result = subprocess.run([sys.executable, '-c', script], env={**os.environ, 'PATH': f'{bin_dir}:{os.environ["PATH"]}'}, stderr=subprocess.PIPE, text=True) assert result.returncode == 0, result.stderr - assert '1 MLOPart devices' in result.stderr, result.stderr + assert 'Traceback' not in result.stderr, result.stderr + assert 'MLOPartUnavailable' not in result.stderr, result.stderr print(' > detected, `localize` refused, SM table available') print() From 84a13a481415dca86cc413c119cc71cf0806dc78 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Mon, 12 Oct 2026 21:46:40 +0000 Subject: [PATCH 3/3] Use the active package name for the epilogue module Signed-off-by: Chris Leonard --- deep_gemm/_C.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deep_gemm/_C.py b/deep_gemm/_C.py index bb909bb785..fd5eddaea5 100644 --- a/deep_gemm/_C.py +++ b/deep_gemm/_C.py @@ -628,7 +628,7 @@ def slice_input_buffers(buffer): # GEMM epilogue classes for the `epilogue` argument, as `deep_gemm.epilogue.X` # NOTES: `alpha` and the FP8 `(d, sfd)` output pair are the shorthand for `Alpha` and `FP8Quantization` -epilogue = types.ModuleType('deep_gemm.epilogue', 'GEMM epilogue classes') +epilogue = types.ModuleType(f'{__package__}.epilogue', 'GEMM epilogue classes') epilogue.Identity = lambda: _torch_ops.epilogue_identity() epilogue.Alpha = lambda alpha: _torch_ops.epilogue_alpha(alpha) epilogue.FP8Quantization = lambda sfd: _torch_ops.epilogue_fp8_quantization(sfd)