[cleonard530]

Summary

Migrate DeepGEMM off of the unstable libtorch ABI and onto the torch Stable ABI. Part of the vllm stable abi migration planned outlined in vllm-project/vllm#26946.

This is "stacked" on top of the PR to migrate DeepGEMM off of pybind and onto TORCH_LIBRARY here #15.

Design choices

A few methods did not have a direct stable equivelent method, so a few wrapper methods were defined in csrc/utils/torch_compat.hpp. Some of these are just wrappers around torch_call_dispatcher (a method that allows us to call the aten methods directly through the dispatcher, so no need for external unstable binary symbols), some had to have their input types converted (e.g. from an int to a Tensor), and some where just added to because the current stable shims could not handle const dtypes yet. There is a short comment above each wrapper describing why it is there.

torch::arange was handled as a special case since there is no stable equivalent and it could not be wrapped by torch_call_dispatcher. A new method, deep_gemm_arange_int64, in ‎csrc/jit_kernels/impls/smxx_layout.hpp was added to provide a fast arange function for deep_gemm. This is compiled using DeepJIT so will be slower on the first run than torch::arange, and is used in the transform_sf_into_required_layout method in ‎csrc/apis/layout.hpp. A stable wrapper using repeat_interleave was considered but showed to have a non-negligible slowdown compared to this custom kernel.

When DeepGEMM still registered it's ops with pybind, get_jit returned an empty Runtime object backed by a std::shared_ptr<Runtime<Backend>. No C++ Runtime methods were exposed as Python attributes, so it acted as an opaque handle. This method was replaced with a get_jit method that calls jit.get() like it did before (through _ensure_jit) and then wraps the output with a blank Runtime class in python (see DeepJIT Runtime comment2). This class was added to preserve API names to support checks such as isinstance(_C.get_jit(), _C.Runtime) but can be removed if not used anywhere. This returned object lost the ability for the method to be used by another binding which accepts a std::shared_ptr<Runtime<Backend>>. Searching open repos on GitHub, I could not find any uses of this object this way (or any calls to this method at all).

For get_bf16_mega_gate_config, STABLE_TORCH_LIBRARY cannot register a dict, so a list was registered instead and the Python compatibility shim (deep_gemm/_C.py) wraps the dict around the list values so that python end users can still use it the same way.

In a few locations, the unstable c10::cuda::CUDAStream was replaced with a std::pair<int32_t, cudaStream_t> (StreamKey) since the device info was needed as well as the stream. Since StreamKey is not hashable, std::unordered_map<c10::cuda::CUDAStream... were replaced with std::map<torch_compat::StreamKey.... These are small maps so the difference in runtime should be negligible.

In csrc/apis/attention.hpp, we used to have DG_HOST_ASSERT(c10::cuda::currentStreamCaptureStatusMayInitCtx() == c10::cuda::CaptureStatus::None); which doesn't have a stable equivalent. It was replaced with not torch_compat::is_capturing(stream.second) which checks capture status on the exact CUDA stream associated with the reference tensor’s device, instead of PyTorch’s current stream.

csrc/apis/epilogue_class.hpp‎ used to register an m.class_<Epilogue>("Epilogue"); and multiple methods took in epilogue_class::EpilogueArg& epilogue as an argument. STABLE_TORCH_LIBRARY can't register classes, so the wrapper to register these ops now takes in  int64_t& epilogue_kind, std::optional<double>& epilogue_alpha, and std::optional<torch::stable::Tensor>& epilogue_sfd and passes them into epilogue_class::unwrap(epilogue_kind, epilogue_alpha, epilogue_sfd) (e.g. fp8_fp4_gemm_nt in gemm.hpp around line 922). Before, in python epilogue.Identity, epilogue.Alpha, epilogue.FP8Quantization, and epilogue.BF16StochasticRounding were registered through TORCH_LIBRARY with something like "epilogue_identity() -> __torch__.torch.classes.deep_gemm.Epilogue". Now, the same calls create a python object _EpilogueSpec, the value is unwrapped in the python method _unpack_epilogue, and is passed into the torch.ops. It's more complicated this way, but gets us around the python torch.ops barrier with only allowed stable types, and since there is no more dispatch overhead, it's slightly faster than before (I think the runtime is still basically noise either way, so IDK how much that matters). The final C++ behavior of the Epilogue classes are unchanged. What changed is its transport to get there.

You can also see these changes in the comments below

Test

pytest tests/test_bf16.py
pytest tests/test_fp8_fp4.py
pytest tests/test_layout.py
pytest tests/test_einsum.py
pytest tests/test_attention.py
pytest tests/test_hyperconnection.py
pytest tests/test_legacy.py
pytest test_torch_library.py
pytest tests/test_coverage_gaps.py  # (local test that fills in the kernel gaps)

Summary: 9/9 passed
RESULT: PASSED

Test that require Blackwell GPUs have not been tested yet. These include

test_mega_gate.py
test_mega_mhc.py
test_mega_moe.py
All six test_sm120_*.py files

and a few test in the other files.

test_layout.py::test_k_grouped_sf_layout_kernels
test_layout.py::test_k_grouped_psum_sf_layout_kernels
test_einsum.py::test_fp8_bhd_hdr_bhr
test_einsum.py::test_fp8_bhd_bhr_hdr
test_attention.py::test_sparse_mqa_logits

---

Torch abi audit (torch-abi-audit) of a vLLM build with this commit:

 -- extensions --
    [STABLE  ] [abi3-ok               ] _C_stable_libtorch.abi3.so  (stable_shim=88, unstable=0)
    [STABLE  ] [abi3-ok               ] _deepselect_C.abi3.so  (stable_shim=40, unstable=0)
    [STABLE  ] [abi3-ok               ] _flashkda_C.abi3.so  (stable_shim=43, unstable=0)
    [STABLE  ] [abi3-ok               ] _flashmla_C.abi3.so  (stable_shim=66, unstable=0)
    [STABLE  ] [abi3-ok               ] _flashmla_extension_C.abi3.so  (stable_shim=63, unstable=0)
    [STABLE  ] [abi3-ok               ] _moe_C_stable_libtorch.abi3.so  (stable_shim=72, unstable=0)
    [STABLE  ] [abi3-ok               ] _qutlass_C.abi3.so  (stable_shim=63, unstable=0)
    [NO-TORCH] [abi3-ok               ] cumem_allocator.abi3.so
    [NO-TORCH] [abi3-ok               ] fs_io_C.abi3.so
    [NO-TORCH] [abi3-ok               ] spinloop.abi3.so
    [STABLE  ] [abi3-ok               ] third_party/deep_gemm/_C_extension.abi3.so  (stable_shim=78, unstable=0)
    [STABLE  ] [abi3-ok               ] vllm_flash_attn/_vllm_fa2_C.abi3.so  (stable_shim=69, unstable=0)
    [STABLE  ] [abi3-ok               ] vllm_flash_attn/_vllm_fa3_C.abi3.so  (stable_shim=63, unstable=0)

Making vLLM completely stable!!!

Performance Benchmarks

Link to performance benchmarks: cleonard530/Benchmarks#1
The tldr; most gemm kernels stayed the same except the fp8_fp4_gemm_{nn,nt,tn,tt} kernels which had a 10-20% runtime increase (2-6 µs/call). I ran another workload with them in a vLLM-like workload. In eager mode the regression was about the same, but it disappeared with the CUDA-graph run. the runtime increase in fp8_fp4_gemm_{nn,nt,tt} went away when the data element_size was cached. fp8_fp4_gemm_tn still has a slight runtime regression.

[cleonard530]

JitRuntimeHandle was removed because the PyTorch stable ABI does not support registering custom C++ classes. The Python Runtime wrapper preserves the existing get_jit() API while _ensure_jit() still initializes DeepJIT; no runtime methods are lost because DeepJIT’s Python binding did not expose any.

[cleonard530]

One qualification: this preserves the API names and initialization side effect, but not the original handle semantics. Previously, get_jit() returned an object backed by std::shared_ptr<Runtime<Backend>>, which could be retained or passed to another compatible C++ binding. The replacement Runtime object does not contain that C++ handle. I found no current DeepGEMM usage of this behavior, but internal downstream users could potentially depend on it.

[janeyx99]

I think what we should document is the highlevel impact of this change. What behavior is held constant and what is now different? This would give maintainers a better picture of how to decide. Also, is it even better to "spoof" the old behavior? We should be as clear as possible about what is changing, and not do extra work to look the same if we cannot promise identical behavior.

[cleonard530]

The wrapper to preserve the existing get_jit() API.

The Python Runtime placeholder preserves the legacy API’s surface semantics, but not the original Torch custom-class type identity. It is intentionally empty because DeepJIT exposed no runtime methods through Python.

[cleonard530]

PyTorch does not have a stable ABI equivalent of at::cuda::getCurrentCUDABlasLtHandle() so we remove this option. I have not been able to find any test that use this and by default it is not used, so I don't know of any major impact, but this should be looked at closely.

[janeyx99]

Huh I think we should be able to undo this deletion:

Codex says:

torch_get_current_cuda_blas_handle (torch/csrc/stable/c/shim.h:126) exists. The unstable C++ at::cuda::getCurrentCUDABlasLtHandle() does exist, but on CUDA it obtains the ordinary managed cuBLAS handle and casts it.

Therefore DeepGEMM can safely do:

void* handle = nullptr;
TORCH_ERROR_CODE_CHECK(torch_get_current_cuda_blas_handle(&handle));
return reinterpret_cast<cublasLtHandle_t>(handle);

[janeyx99]

This should be fixed now!

[cleonard530]

There is no stable shim for torch arange, so we implemented a deep_gemm kernel for this. This now uses DeepJIT for compilation on the first use.

[janeyx99]

We could dispatch to repeat_interleave.Tensor and it would be more tested/stable, but come at the cost of performance after warmup.

[cleonard530]

c10::Dict<std::string, int64_t> cannot be used in a stable operator schema. The stable ABI supports int[], so the C++ implementation returns std::vector<int64_t>. The Python wrapper converts that vector back into the public dictionary

[cleonard530]

The new not torch_compat::is_capturing(stream.second) check is more explicit and stable-ABI compatible: it checks capture status on the exact CUDA stream associated with the reference tensor’s device. The previous c10::cuda::currentStreamCaptureStatusMayInitCtx() == c10::cuda::CaptureStatus::None call used an unstable PyTorch API and queried PyTorch’s current stream instead.