vllm - 💡(How to fix) Fix [Build] DeepGEMM wheel integration: planned cleanups

🚀 The feature, motivation and pitch

Tracking issue for cleanups to the vendored DeepGEMM build path. Companion to upstream DeepGEMM's issue https://github.com/deepseek-ai/DeepGEMM/issues/333

Current state

vLLM bundles DeepGEMM under vllm/third_party/deep_gemm/ and ships a _C.cpython-X.Y-*.so for every CPython covered by requires-python, all side-by-side in the wheel. Python's loader picks the right one at import. The build is in cmake/external_projects/deepgemm.cmake and tools/build_deepgemm_C.py.

This shape was settled on after the single-_C.abi3.so approach was tried and reverted (see #41476 / #41512); landed in #41516.

What we'd like to clean up

1. Collapse the per-Python _C build down to a single _C.abi3.so.

Today's per-Python machinery exists because DeepGEMM's binding uses PYBIND11_MODULE, which inevitably links private CPython symbols (_PyObject_GetDictPtr, _PyThreadState_UncheckedGet, etc.) — so the .so is Python-version-coupled regardless of file naming.

Two paths to fix:

• Upstream switches to TORCH_LIBRARY + a pytorch-shim layer (cf. [vllm-flash-attention/csrc/common/pytorch_shim.h](https://github.com/vllm-project/flash-attention/blob/main/csrc/common/pytorch_shim.h)). Preferred. Coordinate with upstream.
• Maintain a binding fork in vLLM that includes DeepGEMM's csrc/apis/*.hpp and registers ops via TORCH_LIBRARY. Avoids upstream coordination at the cost of mirroring binding changes when DeepGEMM adds ops.

Either path eliminates tools/build_deepgemm_C.py, tools/setup_deepgemm_pythons.sh, tools/check_wheel_deepgemm.py, the per-Python foreach in deepgemm.cmake, and the DEEPGEMM_PYTHON_INTERPRETERS plumbing in the Dockerfile.

2. AOT-compile DeepGEMM's CUDA kernels.

DeepGEMM currently JIT-compiles its kernels at runtime, which is why we vendor CUTLASS + CCCL + DeepGEMM headers under vllm/third_party/deep_gemm/include/. That's several MB of wheel size and a hard runtime dependency on a CUDA toolkit being present on the host.

Switching to AOT compilation would let us drop the vendored headers and the toolkit-at-runtime requirement. Requires upstream changes.

Alternatives

No response

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.