Skip to content

Fix CuTeDSL editable install setup script - #3486

Open
VaggelisGian wants to merge 2 commits into
NVIDIA:mainfrom
VaggelisGian:fix-dsl-prep-editable-libs-wheels
Open

Fix CuTeDSL editable install setup script#3486
VaggelisGian wants to merge 2 commits into
NVIDIA:mainfrom
VaggelisGian:fix-dsl-prep-editable-libs-wheels

Conversation

@VaggelisGian

Copy link
Copy Markdown

Fixes #3416

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 ("0 lib files, 0 Python files copied") and a following pip install -e . failed with ModuleNotFoundError: No module named 'cutlass._mlir'.

The script now downloads the wheels that actually hold the pieces:

  • nvidia-cutlass-dsl-libs-base==<release> provides the cutlass._mlir python package (dialect bindings glue),
  • nvidia-cutlass-dsl-libs-cu12==<release> (or -cu13 via a new --cu13 flag, matching setup.sh conventions) provides the toolkit-native _cutlass_ir binding and libcute_dsl_runtime.so,
  • both are pinned to the release version parsed from the main wheel, whose pin itself is taken from this tree's requirements.txt.

It copies _mlir into the local cutlass/ package (base glue merged with the toolkit-native files), bridges this tree's own _mlir_helpers into cutlass/_mlir_helpers where its relative imports expect them, places runtime libraries under cuXX/lib mirroring an installed wheel so the DSL library discovery finds them, removes a stale flat lib/ 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/.gitignore for the generated paths so running the script does not dirty git status.

Test Plan:

$ cd python/CuTeDSL && python prep_editable_install.py
INFO: Summary: 65 _mlir files, 9 _mlir_helpers files, 3 runtime library files (cu12/lib)
$ python prep_editable_install.py --cu13
INFO: Summary: 65 _mlir files, 9 _mlir_helpers files, 3 runtime library files (cu13/lib)
$ python -m py_compile prep_editable_install.py    # clean

# sha256 of every extracted file compared against the wheel members: 65+9+3 files, zero mismatches

$ docker run --rm -v <repo>:/src python:3.12-slim bash -c \
    "cd /src/python/CuTeDSL && python prep_editable_install.py && \
     pip install numpy cuda-python typing_extensions && pip install -e . && \
     python -c 'import cutlass; import cutlass.cute as cute; print(\"ok\")'"
import ok: /src/python/CuTeDSL/cutlass/__init__.py
cute ok

Before this change the same flow ends at ModuleNotFoundError: No module named 'cutlass._mlir'.

Note: numpy, cuda-python and typing_extensions are needed at import time but pyproject.toml declares 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.

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.
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.

[BUG] CuTeDSL editable install broken: prep_editable_install.py extracts nothing from the split nvidia-cutlass-dsl wheels (>= 4.4.0)

1 participant