Add PyTorch stable ABI extension for libtorch compatibility#2813
Draft
pstjohn wants to merge 3 commits intoNVIDIA:mainfrom
Draft
Add PyTorch stable ABI extension for libtorch compatibility#2813pstjohn wants to merge 3 commits intoNVIDIA:mainfrom
pstjohn wants to merge 3 commits intoNVIDIA:mainfrom
Conversation
Member
|
Why would we need to have a separate stable directory for those files? If it works, we would want to use it by default, no? Also, then maybe the diff would be easier to understand and review? |
pstjohn
force-pushed
the
pstjohn/libtorch-stable-abi-clean
branch
2 times, most recently
from
c7d2553 to
64a7124
Compare
Contributor
Author
I think it's going to be a large diff, but I'll definitely try to cut it back once i have all the tests passing locally |
pstjohn
force-pushed
the
pstjohn/libtorch-stable-abi-clean
branch
8 times, most recently
from
5bd5b78 to
3c8b870
Compare
pstjohn
commented
Mar 31, 2026
pstjohn
force-pushed
the
pstjohn/libtorch-stable-abi-clean
branch
4 times, most recently
from
8a768f4 to
6e19507
Compare
Replace the pybind11-based `transformer_engine_torch` extension with a pure C++ stable ABI extension (`te_stable_abi`) that uses PyTorch's `torch::Library` API. This eliminates the dependency on unstable PyTorch C++ internals (ATen, c10, pybind11), making TE compatible across PyTorch versions without recompilation. Key changes: - New `_stable_torch_module.py` routes all ops through stable ABI - New `_quantize_stable.py` handles FP8/NVFP4/MXFP8 quantization - C++ extensions (attention, gemm, permutation, etc.) ported to stable ABI tensor wrappers - CMakeLists.txt handles suffixed CUDA archs (100a, 103a, etc.) - Float8Quantizer computes transpose after quantization - FP8 attention backward properly feeds amax to global state - GELU epilogue fusion in generic_gemm - CUDA graph capture compatible with FP8 quantization Signed-off-by: Peter St. John <pstjohn@nvidia.com>
pstjohn
force-pushed
the
pstjohn/libtorch-stable-abi-clean
branch
from
6e19507 to
e48c15c
Compare
Contributor
Author
|
/te-ci pytorch |
- Add C++ group_quantize op wrapping nvte_group_quantize for fused MXFP8 quantization instead of per-chunk Python loop - Fix segfault in grouped GEMM discrete_out by keeping D metadata tensors on CPU (C++ dereferences pointers on host) - Fix FP8 attention NaN by skipping scale_inv initialization for CurrentScaling quantizers (scale is uninitialized, kernel computes on-the-fly) - Fix NVFP4 backward by routing through two-step dactivation path (fused kernel asserts IS_DACT not supported for NVFP4) - Fix scalar tensor columnwise dequant in quantize_into for 0-dim tensors - Fix MXFP8 columnwise dequant by transposing data+scale before kernel call - Fix CUDA graph capture by caching first_dims.tolist() results - Fix swizzle_scale_for_gemm C++ schema to include columnwise param - Add post-quantize swizzle in MXFP8Quantizer for optimize_for_gemm - Clean up build_tools/pytorch.py: restore NVSHMEM handling, remove unnecessary absolute rpath, minimize diff vs main - Update MXFP8 test assertion for preserved rowwise data Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Documents pure-Python fallbacks in _stable_torch_module.py that replace fused C++ kernels, organized by severity. Serves as a roadmap for porting remaining operations to the stable ABI C++ extension. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Implement a stable ABI layer that replaces the pybind11-based C++ extension with torch::Library-registered operations using torch::stable::Tensor. This allows the PyTorch extension to be built once and work across multiple Python/PyTorch versions without recompilation.
Key changes: