Fix CuTeDSL editable install setup script - #3486
Open
VaggelisGian wants to merge 2 commits into
Open
Conversation
prep_editable_install.py downloaded only the pure-Python
nvidia-cutlass-dsl wheel and globbed it for lib/*.so. Since the wheel
split, that wheel carries no files at all, so the script copied nothing
and a following pip install -e . failed with
"No module named 'cutlass._mlir'".
Download the wheels that now hold the pieces: nvidia-cutlass-dsl-libs-base
for the cutlass._mlir python package and nvidia-cutlass-dsl-libs-cu12 or
cu13 (--cu13) for the toolkit-native _cutlass_ir binding, both pinned to
the release version of the main wheel from requirements.txt. Copy the
tree's own _mlir_helpers into the cutlass package where its relative
imports expect it, place runtime libraries under cuXX/lib like an
installed wheel does, remove a stale flat lib/ left by older runs that
would shadow the new layout, and fail loudly when no native binding was
extracted instead of reporting success.
Add python/CuTeDSL/.gitignore for the generated paths so running the
script does not dirty git status.
Test Plan:
python prep_editable_install.py on Windows:
Summary: 65 _mlir files, 9 _mlir_helpers files,
3 runtime library files (cu12/lib)
python prep_editable_install.py --cu13:
same summary with cu13/lib and the _cutlass_ir.cu13 binding
python -m py_compile prep_editable_install.py: clean
docker python:3.12-slim: prep script + pip install -e . plus
numpy, cuda-python, typing_extensions:
import ok: /src/python/CuTeDSL/cutlass/__init__.py
cute ok
Old script (before this change): copied 0 files; pip install -e .
then fails with ModuleNotFoundError: No module named 'cutlass._mlir'
copy_runtime_libs only warned when the libs wheel carried no cuXX/lib directory, so a future upstream layout change would produce an install that imports fine and then fails confusingly at first JIT compile. The missing-native guard already hard-fails for _cutlass_ir; extend the same treatment to the runtime libraries. Test Plan: python:3.12-slim container against today's 4.7.0 stub+libs wheels: prep completes with '3 runtime library files (cu12/lib)' summary, unchanged happy path. Guard triggers only when lib_files_copied == 0.
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
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.
Fixes #3416
prep_editable_install.pydownloaded only the pure-Pythonnvidia-cutlass-dslwheel and globbed it forlib/*.so. Since the wheel split, that wheel carries no files at all, so the script copied nothing ("0 lib files, 0 Python files copied") and a followingpip install -e .failed withModuleNotFoundError: No module named 'cutlass._mlir'.The script now downloads the wheels that actually hold the pieces:
nvidia-cutlass-dsl-libs-base==<release>provides thecutlass._mlirpython package (dialect bindings glue),nvidia-cutlass-dsl-libs-cu12==<release>(or-cu13via a new--cu13flag, matchingsetup.shconventions) provides the toolkit-native_cutlass_irbinding andlibcute_dsl_runtime.so,requirements.txt.It copies
_mlirinto the localcutlass/package (base glue merged with the toolkit-native files), bridges this tree's own_mlir_helpersintocutlass/_mlir_helperswhere its relative imports expect them, places runtime libraries undercuXX/libmirroring an installed wheel so the DSL library discovery finds them, removes a stale flatlib/left behind by older runs that would shadow the new layout, and fails loudly if no native binding was extracted instead of reporting success.Adds
python/CuTeDSL/.gitignorefor the generated paths so running the script does not dirty git status.Test Plan:
Before this change the same flow ends at
ModuleNotFoundError: No module named 'cutlass._mlir'.Note:
numpy,cuda-pythonandtyping_extensionsare needed at import time butpyproject.tomldeclares no dependencies; that pre-existing gap is untouched here. The editable install remains Linux-only at runtime because the extracted natives are ELF shared objects.