From 3042dd02057491404a6b3c0d19e29fcbc6047b05 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Mon, 12 Oct 2026 00:22:17 +0000 Subject: [PATCH] Register the locality-domain shim under the _C module's own name When DeepGEMM is vendored under another package name (vLLM ships it as vllm.third_party.deep_gemm), the hardcoded 'deep_gemm._C.locality_domain' key does not match the relative 'from ._C.locality_domain import' in locality_domain.py, so importing the package fails. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Yongye Zhu --- deep_gemm/_C.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/deep_gemm/_C.py b/deep_gemm/_C.py index c108f3e4da..05677ad635 100644 --- a/deep_gemm/_C.py +++ b/deep_gemm/_C.py @@ -632,7 +632,7 @@ def slice_input_buffers(buffer): # Memory homed in the GPU's locality domains, re-exported by `deep_gemm.locality_domain` -locality_domain = types.ModuleType('deep_gemm._C.locality_domain', "Memory homed in the GPU's locality domains") +locality_domain = types.ModuleType(f'{__name__}.locality_domain', "Memory homed in the GPU's locality domains") locality_domain.get_num_locality_domains = _torch_ops.locality_domain_get_num_locality_domains locality_domain.get_granularity = _torch_ops.locality_domain_get_granularity locality_domain.is_localization_available = _torch_ops.locality_domain_is_localization_available @@ -642,7 +642,8 @@ def slice_input_buffers(buffer): locality_domain.get_sm_locality_domains = _torch_ops.locality_domain_get_sm_locality_domains locality_domain.get_balanced_sm_locality_domains = _torch_ops.locality_domain_get_balanced_sm_locality_domains locality_domain.release_mlopart = _torch_ops.locality_domain_release_mlopart -# NOTES: registered so that `from ._C.locality_domain import ...` resolves as with the former pybind submodule +# NOTES: registered so that `from ._C.locality_domain import ...` resolves as with the former pybind submodule, +# named after this module so the import also resolves when the package is vendored under another name sys.modules[locality_domain.__name__] = locality_domain