Skip to content

Register the locality-domain shim under the _C module's own name - #31

Merged
zyongye merged 1 commit into
devfrom
fix/vendored-locality-domain
Oct 12, 2026
Merged

zyongye merged 1 commit into
devfrom
fix/vendored-locality-domain

Conversation

@zyongye

@zyongye zyongye commented Oct 12, 2026

Copy link
Copy Markdown
Member

#30 made _C.py register the locality-domain shim as sys.modules['deep_gemm._C.locality_domain'] so that locality_domain.py can do from ._C.locality_domain import .... The key is hardcoded, but the relative import resolves against the package's real name. vLLM vendors this package as vllm.third_party.deep_gemm, so the lookup is vllm.third_party.deep_gemm._C.locality_domain, which is never registered, and import vllm.third_party.deep_gemm fails:

ModuleNotFoundError: No module named 'vllm.third_party.deep_gemm._C.locality_domain'; 'vllm.third_party.deep_gemm._C' is not a package

vLLM treats an import failure as "DeepGEMM unavailable", so this blocks bumping vLLM's pin to dev (15c02c8).

The fix names the shim after the _C module itself (f'{__name__}.locality_domain'). A top-level install still registers deep_gemm._C.locality_domain, so its behaviour is unchanged.

Testing

On GB200 with the 15c02c8 build plus this _C.py:

  • import vllm.third_party.deep_gemm succeeds, and deep_gemm.get_num_locality_domains() returns 2.
  • A top-level import deep_gemm.locality_domain succeeds, and the shim is still named deep_gemm._C.locality_domain.

Not covered

The MLOPart bridge in csrc/runtime/mlopart.hpp also hardcodes PyImport_ImportModule("deep_gemm.utils.mlopart"). That path only runs on explicit localize() calls, which vLLM does not make, so it's left for a follow-up.

🤖 Generated with Claude Code

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) <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
@zyongye
zyongye force-pushed the fix/vendored-locality-domain branch from fa9dfe2 to 3042dd0 Compare October 12, 2026 03:20
@zyongye
zyongye merged commit 1448881 into dev Oct 12, 2026
1 check passed
zyongye added a commit to zyongye/vllm that referenced this pull request Oct 13, 2026
…parse indexer

Pin vllm-project/DeepGEMM dev 6ce1c5e: the upstream 26/09/30 release
(vllm-project/DeepGEMM#30), the vendoring import fixes
vllm-project/DeepGEMM#31 and vllm-project/DeepGEMM#32, and the paged
sparse MQA page-stride relaxation vllm-project/DeepGEMM#33. Vendor the
new deep_gemm/locality_domain.py. DeepGEMM#32 makes the MLOPart bridge
and the epilogue module use the package's own name, so they also work
when vendored as vllm.third_party.deep_gemm. DeepGEMM#33 lowers the
page-stride guard from 512 to 256 bytes; the DeepSeek-V4.1
sparse-logits block stride stays 512-aligned, which still satisfies it.

The MQA logits API changed: clean_logits/logits_dtype are gone,
max_seqlen_k is required, prefill logits are compressed to each row's
K span, and SM100 takes only MXFP8/MXFP4 with BF16 weights and returns
BF16 logits. On SM100 the FP8 indexer K writers store their existing
power-of-two scale as four packed UE8M0 bytes (a lossless MXFP8
encoding), FP8 Q passes unit MX scales with its scale still folded into
the weights, and BF16 logits are upcast row-aware for the FP32 top-k
kernels. check_wheel_deepgemm.py now imports the vendored package.

build_deepgemm_C.py passes --force to setup.py build_ext: DeepGEMM's
setup.py only tracks python_api.cpp, so an incremental build after a
header-only pin change kept a stale extension. CMake already decides
when this step reruns.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant