Skip to content

Port csrc extension to PyTorch stable ABI#308

Open
tiran wants to merge 1 commit into
ovg-project:mainfrom
tiran:torch-stable-abi
Open

Port csrc extension to PyTorch stable ABI#308
tiran wants to merge 1 commit into
ovg-project:mainfrom
tiran:torch-stable-abi

Conversation

@tiran
Copy link
Copy Markdown

@tiran tiran commented Apr 17, 2026

Replace pybind11 + torch/extension.h bindings with PyTorch's stable ABI (STABLE_TORCH_LIBRARY, torch::stable::Tensor, torch::stable::Device). This decouples the extension from specific PyTorch C++ ABI versions.

  • Register ops via STABLE_TORCH_LIBRARY + TORCH_BOX instead of PYBIND11_MODULE
  • Replace torch::Tensor/Device/Dtype with stable::Tensor/Device and headeronly::ScalarType throughout
  • Replace torch::from_blob with torch::stable::from_blob
  • Add kvcached/vmm_ops.py wrapper that loads _vmm_ops_lib and re-exports ops via torch.ops.kvcached.*, preserving existing import API
  • Remove torch_python from linked libraries, rename extension to _vmm_ops_lib
  • Remove dead free_ftensor_() and torch_utils.ipp

Fixes: #306

Replace pybind11 + torch/extension.h bindings with PyTorch's stable ABI
(STABLE_TORCH_LIBRARY, torch::stable::Tensor, torch::stable::Device).
This decouples the extension from specific PyTorch C++ ABI versions.

- Register ops via STABLE_TORCH_LIBRARY + TORCH_BOX instead of PYBIND11_MODULE
- Replace torch::Tensor/Device/Dtype with stable::Tensor/Device and
  headeronly::ScalarType throughout
- Replace torch::from_blob with torch::stable::from_blob
- Add kvcached/vmm_ops.py wrapper that loads _vmm_ops_lib and re-exports
  ops via torch.ops.kvcached.*, preserving existing import API
- Remove torch_python from linked libraries, rename extension to _vmm_ops_lib
- Remove dead free_ftensor_() and torch_utils.ipp

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: Christian Heimes <cheimes@redhat.com>
@tiran tiran force-pushed the torch-stable-abi branch from 2252a03 to d077266 Compare April 20, 2026 07:40
@tiran
Copy link
Copy Markdown
Author

tiran commented Apr 20, 2026

I have addressed the code format issues with pre-commit run -a clang-format.

@tiran tiran marked this pull request as ready for review April 20, 2026 07:41
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.

[RFE] Port kvcached to libtorch stable ABI

1 participant