Skip to content

migrated DeepGEMM to torch stable ABI (unchecked) #8

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
cleonard530 wants to merge 5 commits into
base: migrate_pybind_to_torch_library_dev
Choose a base branch
from
Open

migrated DeepGEMM to torch stable ABI (unchecked) #8

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
286f765
migrated DeepGEMM to torch stable ABI (unchecked)
cleonard530 Sep 18, 2026
b06ddc9
moved min torch version from 2.13 to 2.10
cleonard530 Sep 28, 2026
ffa3a29
implemented arange kernel for deep_gemm since torch does not have a s…
cleonard530 Sep 29, 2026
4bba562
Clean up stable ABI compatibility wrappers and strided tensor allocat…
cleonard530 Sep 30, 2026
059e168
Cache GEMM element sizes in descriptors and reuse them across SM90, S…
cleonard530 Oct 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 12 additions & 27 deletions .github/workflows/_build.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,11 +16,7 @@ on:
required: true
type: string
torch-version:
description: "The PyTorch version to use for the build"
required: true
type: string
cxx11_abi:
description: "The C++11 ABI to use for the build"
description: "PyTorch build version (2.10 or newer)"
required: true
type: string
upload-to-release:
Expand All @@ -45,7 +41,7 @@ defaults:
jobs:
build-wheel:
runs-on: ${{ inputs.runs-on }}
name: Build wheel (${{ inputs.release-version }}-${{ inputs.python-version }}-${{ inputs.cuda-version }}-${{ inputs.torch-version }}-${{ inputs.cxx11_abi }})
name: Build wheel (${{ inputs.release-version }}-${{ inputs.python-version }}-${{ inputs.cuda-version }}-${{ inputs.torch-version }})
steps:
- name: Checkout
uses: actions/checkout@v4
Expand Down Expand Up @@ -101,28 +97,17 @@ jobs:
- name: Install PyTorch ${{ inputs.torch-version }}+cu${{ inputs.cuda-version }}
run: |
pip install --upgrade pip
# With python 3.13 and torch 2.5.1, unless we update typing-extensions, we get error
# AttributeError: attribute '__default__' of 'typing.ParamSpec' objects is not writable
pip install typing-extensions==4.12.2
# We want to figure out the CUDA version to download pytorch
# e.g. we can have system CUDA version being 11.7 but if torch==1.12 then we need to download the wheel from cu116
# see https://github.com/pytorch/pytorch/blob/main/RELEASE.md#release-compatibility-matrix
# This code is ugly, maybe there's a better way to do this.
export TORCH_CUDA_VERSION=$(python -c "from os import environ as env; \
minv = {'2.4': 118, '2.5': 118, '2.6': 118, '2.7': 118, '2.8': 126}[env['MATRIX_TORCH_VERSION']]; \
maxv = {'2.4': 124, '2.5': 124, '2.6': 126, '2.7': 128, '2.8': 129}[env['MATRIX_TORCH_VERSION']]; \
print(minv if int(env['MATRIX_CUDA_VERSION']) < 120 else maxv)" \
)
# The stable target needs torch >= 2.10. Use the selected CUDA index.
export TORCH_CUDA_VERSION=$MATRIX_CUDA_VERSION
if [[ ${{ inputs.torch-version }} == *"dev"* ]]; then
# pip install --no-cache-dir --pre torch==${{ inputs.torch-version }} --index-url https://download.pytorch.org/whl/nightly/cu${TORCH_CUDA_VERSION}
# Can't use --no-deps because we need cudnn etc.
# Hard-coding this version of pytorch-triton for torch 2.6.0.dev20241001
pip install jinja2
pip install https://download.pytorch.org/whl/nightly/pytorch_triton-3.1.0%2Bcf34004b8a-cp${MATRIX_PYTHON_VERSION}-cp${MATRIX_PYTHON_VERSION}-linux_x86_64.whl
pip install --no-cache-dir --pre https://download.pytorch.org/whl/nightly/cu${TORCH_CUDA_VERSION}/torch-${{ inputs.torch-version }}%2Bcu${TORCH_CUDA_VERSION}-cp${MATRIX_PYTHON_VERSION}-cp${MATRIX_PYTHON_VERSION}-linux_x86_64.whl
# The previous Torch 2.6 nightly path hard-coded its matching pytorch-triton
# wheel; let pip resolve the matching dependencies for the requested nightly.
pip install --no-cache-dir --pre torch==${{ inputs.torch-version }} --index-url https://download.pytorch.org/whl/nightly/cu${TORCH_CUDA_VERSION}
else
pip install --no-cache-dir torch==${{ inputs.torch-version }} --index-url https://download.pytorch.org/whl/cu${TORCH_CUDA_VERSION}
fi
python -c "import torch; from packaging.version import Version; assert Version(torch.__version__).release[:2] >= (2, 10), 'Stable ABI requires torch >= 2.10'"
nvcc --version
python --version
python -c "import torch; print('PyTorch:', torch.__version__)"
Expand All @@ -133,9 +118,9 @@ jobs:
uses: actions/cache/restore@v4
with:
path: build.tar
key: build-${{ inputs.release-version }}-${{ inputs.python-version }}-${{ inputs.cuda-version }}-${{ inputs.torch-version }}-${{ inputs.cxx11_abi }}-${{ github.run_number }}-${{ github.run_attempt }}
key: stable210-build-${{ inputs.release-version }}-${{ inputs.python-version }}-${{ inputs.cuda-version }}-${{ inputs.torch-version }}-${{ github.run_number }}-${{ github.run_attempt }}
restore-keys: |
build-${{ inputs.release-version }}-${{ inputs.python-version }}-${{ inputs.cuda-version }}-${{ inputs.torch-version }}-${{ inputs.cxx11_abi }}-
stable210-build-${{ inputs.release-version }}-${{ inputs.python-version }}-${{ inputs.cuda-version }}-${{ inputs.torch-version }}-

- name: Unpack build cache
run: |
Expand Down Expand Up @@ -177,7 +162,7 @@ jobs:
timeout 5h python setup.py bdist_wheel --dist-dir=dist || EXIT_CODE=$?

if [ $EXIT_CODE -eq 0 ]; then
tmpname=cu${WHEEL_CUDA_VERSION}torch${MATRIX_TORCH_VERSION}cxx11abi${{ inputs.cxx11_abi }}
tmpname=cu${WHEEL_CUDA_VERSION}.torchstable210
wheel_name=$(ls dist/*whl | xargs -n 1 basename | sed "s/-/+$tmpname-/2")
ls dist/*whl |xargs -I {} mv {} dist/${wheel_name}
echo "wheel_name=${wheel_name}" >> $GITHUB_ENV
Expand All @@ -199,7 +184,7 @@ jobs:
if: always() && steps.build_wheel.outputs.build_exit_code == 124
uses: actions/cache/save@v4
with:
key: build-${{ inputs.release-version }}-${{ inputs.python-version }}-${{ inputs.cuda-version }}-${{ inputs.torch-version }}-${{ inputs.cxx11_abi }}-${{ github.run_number }}-${{ github.run_attempt }}
key: stable210-build-${{ inputs.release-version }}-${{ inputs.python-version }}-${{ inputs.cuda-version }}-${{ inputs.torch-version }}-${{ github.run_number }}-${{ github.run_attempt }}
path: build.tar

- name: Log Built Wheels
Expand Down
7 changes: 1 addition & 6 deletions .github/workflows/build.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,11 +17,7 @@ on:
required: true
type: string
torch-version:
description: "The PyTorch version to use for the build"
required: true
type: string
cxx11_abi:
description: "Enable torch flag C++11 ABI (TRUE/FALSE)"
description: "PyTorch build version (2.10 or newer)"
required: true
type: string
upload-to-release:
Expand All @@ -47,7 +43,6 @@ jobs:
python-version: ${{ inputs.python-version }}
cuda-version: ${{ inputs.cuda-version }}
torch-version: ${{ inputs.torch-version }}
cxx11_abi: ${{ inputs.cxx11_abi }}
upload-to-release: ${{ inputs.upload-to-release }}
release-version: ${{ inputs.release-version }}
use-local-version: ${{ inputs.use-local-version }}
8 changes: 1 addition & 7 deletions .github/workflows/publish.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,20 +44,14 @@ jobs:
# setup.py builds one abi3 wheel with a CPython 3.10 floor, so one Python
# build covers every supported Python version in the release matrix.
python-version: ["3.10"]
torch-version: ["2.4.0", "2.5.1", "2.6.0", "2.7.1", "2.8.0"]
torch-version: ["2.10.0"]
cuda-version: ["12.9.1"]
# We need separate wheels that either uses C++11 ABI (-D_GLIBCXX_USE_CXX11_ABI) or not.
# Pytorch wheels currently don't use it, but nvcr images have Pytorch compiled with C++11 ABI.
# Without this we get import error (undefined symbol: _ZN3c105ErrorC2ENS_14SourceLocationESs)
# when building without C++11 ABI and using it on nvcr images.
cxx11_abi: ["FALSE", "TRUE"]
uses: ./.github/workflows/_build.yml
with:
runs-on: ${{ matrix.os }}
python-version: ${{ matrix.python-version }}
cuda-version: ${{ matrix.cuda-version }}
torch-version: ${{ matrix.torch-version }}
cxx11_abi: ${{ matrix.cxx11_abi }}
release-version: ${{ needs.setup_release.outputs.release-version }}
upload-to-release: true
use-local-version: false
Expand Down
5 changes: 4 additions & 1 deletion CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,10 @@ link_directories(${TORCH_INSTALL_PREFIX}/lib ${CUDA_TOOLKIT_ROOT_DIR}/lib64 ${CU
# The main Python API entrance
add_library(_C_extension MODULE csrc/python_api.cpp)
set_target_properties(_C_extension PROPERTIES PREFIX "")
target_compile_definitions(_C_extension PRIVATE DJ_DISABLE_GIL=1)
target_compile_definitions(_C_extension PRIVATE
DJ_DISABLE_GIL=1
TORCH_TARGET_VERSION=0x020a000000000000
USE_CUDA)
target_link_libraries(_C_extension PRIVATE ${TORCH_LIBRARIES} CUDA::cudart)

# Enable kernel code indexing with CMake-based IDEs
Expand Down
4 changes: 2 additions & 2 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,10 +26,10 @@ Despite its lightweight design, DeepGEMM's performance matches or exceeds expert
### Requirements

- NVIDIA SM90 or SM100 architecture GPU
- Python 3.8 or higher
- Python 3.10 or higher
- Compilers and standard libraries with C++20 `<format>` support
- CUDA Toolkit 12.9 or higher
- PyTorch 2.3 or higher
- PyTorch 2.10 or higher
- CUTLASS 4.0 or higher (could be cloned by Git submodule)

### Development
Expand Down
Loading