From be2a4eb14cd95fdb30dd40d654d7143f046a1a9d Mon Sep 17 00:00:00 2001 From: Jari Kolehmainen Date: Fri, 25 Sep 2026 16:27:26 +0000 Subject: [PATCH] Add LCK CUDA-selectable NVSHMEM extras; retain CUDA 12.9 releases (#1492) ## Summary Make the LCK NVSHMEM dependency selectable from the pip command while keeping the combined Hopper/Blackwell release wheel on **CUDA 12.9.1**. ```bash pip install "liger-cute-kernels[cu12]" # For a locally built CUDA 13 Hopper wheel: pip install "/path/to/liger_cute_kernels--py3-none-linux_x86_64.whl[cu13]" ``` - Move NVSHMEM out of unconditional dependencies into `cu12` and `cu13` extras, each pinned to `3.6.5`. Plain LCK leaves NVSHMEM provisioning to the caller. - Keep the existing digest-pinned CUDA 12.9.1 release image, SM90a/SM100f targets, and `nvidia-nvshmem-cu12` build prerequisite. Explicitly package `".[cu12]"` and validate that both optional dependencies activate only for their corresponding extra. - Fix the Hopper non-RDC cubin command to prepend `-I` to every CUDA include directory. CUDA 13 exposes an additional CCCL include directory; previously NVCC interpreted that bare path as a second input file. - Update installation documentation and missing-runtime guidance. Extras select runtime dependencies, not different compiled binaries. Select exactly one matching the core's build toolkit; the two NVSHMEM distributions share an installation namespace. Source builds still require NVSHMEM to be installed before compilation. ### Why the release stays on CUDA 12.9 The complete Hopper core and non-RDC cubin build with CUDA 13.0.88 after the include-path fix. However, the complete SM100 build fails in three MoE backward variants because `ptxas` rejects mixed `.cta_group::1` and `.cta_group::2` instructions in one kernel. This PR does not remove those variants or attempt a kernel redesign. NVIDIA reports the same newer-toolkit restriction in [CCCL #8297](https://github.com/NVIDIA/cccl/pull/8297). The CUDA 13 extra supports appropriately built source artifacts; it does not claim that the combined release wheel builds with CUDA 13. ## Testing Done - Hardware Type: **4 x NVIDIA H200**, CUDA compiler **13.0.88**, Torch **2.11.0.1+cu130**. - **21 focused build/loader tests passed**, including optional dependency metadata, missing-runtime guidance, and CMake include-list expansion with single/multiple CUDA include paths. - Ruff lint and formatting checks passed for the changed Python files. - Parsed the release workflow's embedded Python, checked the packaging shell syntax, and exercised its optional-dependency validation against metadata from the built wheel. - Built, packaged, and installed the complete Hopper core plus non-RDC cubin. The test wheel was 46.97 MiB and used `linux_x86_64`, not a manylinux claim, because that test image uses glibc 2.38. - Ran six numerical cases: Qwen3-30B-A3B, Qwen3-235B-A22B, and Llama-4-Scout at 1,024/2,048 tokens per GPU, EP=4, seed 42. Compared CUDA 13 default and non-RDC-enabled execution against the existing CUDA 12 wheel on the same GPUs and with the same Python frontend and inputs. - Covered forward output, input gradients, all three expert-weight gradients, and selected routing-weight gradients against Torch FP32. All 36 maximum-error metrics matched the CUDA 12 baseline exactly; the change in the equally weighted average mean error was below 0.000008% in both CUDA 13 modes. This is a single-seed comparison, not a claim of bitwise-identical tensors or exhaustive coverage. - Complete CUDA 13 SM100 compilation was attempted and failed as described above. No combined CUDA 13 release wheel was produced. A fresh full CUDA 12.9 release build was not run for this PR; its build image and native compilation settings remain unchanged apart from the compatible include-list fix. The full suites below were not run; the targeted checks and numerical coverage are listed above. - [ ] run `make test` to ensure correctness - [ ] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence Co-authored-by: jkolehma Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../workflows/publish-liger-cute-kernels.yml | 23 +++++---- README.md | 14 +++++- liger_cute_kernels/README.md | 48 +++++++++++++++++-- liger_cute_kernels/csrc/core/CMakeLists.txt | 5 +- .../liger_cute_kernels/__init__.py | 4 +- .../liger_cute_kernels/tvm_ffi.py | 5 +- liger_cute_kernels/setup.py | 9 ++-- liger_cute_kernels/test/test_cute_build.py | 46 ++++++++++++++++++ .../test/test_tvm_ffi_loader.py | 8 ++++ 9 files changed, 139 insertions(+), 23 deletions(-) diff --git a/.github/workflows/publish-liger-cute-kernels.yml b/.github/workflows/publish-liger-cute-kernels.yml index ca30dd95..7b86d737 100644 --- a/.github/workflows/publish-liger-cute-kernels.yml +++ b/.github/workflows/publish-liger-cute-kernels.yml @@ -71,7 +71,7 @@ jobs: - name: Configure CUDA 12.9 dependencies run: | nvcc --version | tee nvcc-version.txt - grep -q 'release 12.9' nvcc-version.txt + grep -q 'release 12\.9,' nvcc-version.txt test "$(getconf GNU_LIBC_VERSION)" = "glibc 2.35" python - <<'PY' >> "$GITHUB_ENV" import importlib.util @@ -185,7 +185,7 @@ jobs: LIGER_CUTE_VERSION: ${{ steps.version.outputs.version }} run: | cd liger_cute_kernels - python -m pip wheel . \ + python -m pip wheel ".[cu12]" \ --no-deps \ --no-build-isolation \ --wheel-dir dist @@ -238,12 +238,19 @@ jobs: if metadata["Version"] != os.environ["EXPECTED_VERSION"]: raise SystemExit(f"wheel version {metadata['Version']} != {os.environ['EXPECTED_VERSION']}") requirements = [Requirement(value) for value in metadata.get_all("Requires-Dist", [])] - nvshmem = next( - (requirement for requirement in requirements if requirement.name == "nvidia-nvshmem-cu12"), - None, - ) - if nvshmem is None or str(nvshmem.specifier) != "==3.6.5": - raise SystemExit(f"missing pinned NVSHMEM dependency: {requirements}") + nvshmem = [requirement for requirement in requirements if requirement.name.startswith("nvidia-nvshmem-")] + if len(nvshmem) != 2 or not {"cu12", "cu13"} <= set(metadata.get_all("Provides-Extra", [])): + raise SystemExit(f"missing CUDA-selectable NVSHMEM extras: {requirements}") + for extra in ("cu12", "cu13"): + matches = [requirement for requirement in nvshmem if requirement.name == f"nvidia-nvshmem-{extra}"] + if len(matches) != 1: + raise SystemExit(f"expected one NVSHMEM requirement for {extra}: {nvshmem}") + requirement = matches[0] + if str(requirement.specifier) != "==3.6.5" or requirement.marker is None: + raise SystemExit(f"NVSHMEM must be pinned and optional: {requirement}") + for selected in ("", "cu12", "cu13"): + if requirement.marker.evaluate({"extra": selected}) != (selected == extra): + raise SystemExit(f"incorrect NVSHMEM extra selection: {requirement}") print(f"{wheel.name}: {wheel.stat().st_size / 1024**2:.1f} MiB") PY diff --git a/README.md b/README.md index 8248c307..4bc6b32b 100644 --- a/README.md +++ b/README.md @@ -166,9 +166,21 @@ $ pip install liger-kernel To also install the matching native Liger communication kernels: ```bash -$ pip install "liger-kernel[lck]" +$ pip install "liger-kernel[lck]" "liger-cute-kernels[cu12]" ``` +The native LCK release wheel build remains on **CUDA 12.9** (CUDA Toolkit 12.9.1). +The optional `cu12` extra installs `nvidia-nvshmem-cu12==3.6.5`; `cu13` +installs `nvidia-nvshmem-cu13==3.6.5` for CUDA 13 source builds. Plain +`liger-cute-kernels` does not install NVSHMEM. Select only one extra, matching +the native wheel's build toolkit: extras select dependencies, not a different +compiled wheel. CUDA 13 source builds currently work for Hopper, but the +combined release build is blocked by mixed CTA-group instructions in the +Blackwell MoE backward kernel. Previously published artifacts are unchanged. The default +`liger-kernel` installation is unchanged. See the +[native wheel build instructions](liger_cute_kernels/README.md#building-the-native-wheel) +for source builds and compatibility details. + To install the nightly version: ```bash diff --git a/liger_cute_kernels/README.md b/liger_cute_kernels/README.md index 26d5cf1f..a73113de 100644 --- a/liger_cute_kernels/README.md +++ b/liger_cute_kernels/README.md @@ -304,14 +304,18 @@ src/liger_kernel/ops/cute/ ## Prerequisites -- **CUDA toolkit** with `nvcc` and either SM 9.0a (Hopper / `sm_90a`) or +- **CUDA toolkit 12.9** with `nvcc` (release builds pin 12.9.1) and either SM 9.0a (Hopper / `sm_90a`) or Blackwell family (`sm_100f`) support. The family target covers both B200 (`sm_100`) and B300 (`sm_103`) while enabling TCGEN05 UMMA and TMEM. + CUDA 13.0 also builds the Hopper path, but currently rejects mixed + `.cta_group::1`/`.cta_group::2` instructions in the Blackwell MoE backward + kernel; the combined release wheel therefore remains on CUDA 12.9. - **NVSHMEM** install (host `.so`, device `.a`, headers). Two layouts are supported: - Native/system install: point `NVSHMEM_HOME` at it, or use the default `/usr/local/nvshmem`. - - PyPI install: install `nvidia-nvshmem-cu12==3.6.5`. The Python wheel + - PyPI install: install `nvidia-nvshmem-cu12==3.6.5` for CUDA 12.9 or + `nvidia-nvshmem-cu13==3.6.5` for CUDA 13. The Python wheel builder and `build_core()` auto-detect that package layout and create unversioned compatibility symlinks for CMake when needed. For direct CMake invocation, pass the package root as `-DNVSHMEM_HOME=...`; the find module @@ -481,8 +485,33 @@ This module's `setup.py` packages the native libraries into the independent **`liger-cute-kernels`** distribution. The release wheel contains both `libliger_cute_kernels_sm90a.so` and `libliger_cute_kernels_sm100f.so`; NVSHMEM is installed independently through -the pinned `nvidia-nvshmem-cu12` dependency. Build against the local -CUDA/NVSHMEM environment (no build isolation), from this module directory: +the optional `cu12` or `cu13` extra, pinned to version `3.6.5`. +The base package does not install either NVSHMEM distribution. The release workflow builds +both cores with CUDA Toolkit **12.9.1** in a digest-pinned +`nvidia/cuda:12.9.1-devel-ubuntu22.04` container. +The Linux x86-64/glibc 2.35 wheel baseline and GPU targets are unchanged. + +Release-wheel consumers need a CUDA 12-compatible runtime and NVIDIA driver. Do not mix +CUDA 12 cores with the CUDA 13 NVSHMEM package or vice versa; both distributions install +into `nvidia/nvshmem`, so use a clean environment rather than installing +`nvidia-nvshmem-cu12` and `nvidia-nvshmem-cu13` together. Previously published +CUDA 12.9 wheels are not changed by this build update. + +Choose the NVSHMEM dependency explicitly when installing a built wheel: + +```bash +python -m pip install "liger-cute-kernels[cu12]" +# For a local Hopper wheel built with CUDA 13: +python -m pip install "/path/to/liger_cute_kernels--py3-none-linux_x86_64.whl[cu13]" +``` + +Extras select dependencies only; they neither choose the CUDA compiler nor +convert the packaged binaries to another CUDA version. Do not request +`[cu12,cu13]`: both packages write to the same namespace. Existing system-managed +NVSHMEM users can install the base wheel without either extra. + +Build against the local CUDA 12.9/NVSHMEM environment (no build isolation), +from this module directory: ```bash cd liger_cute_kernels @@ -493,6 +522,14 @@ LIGER_CUTE_CUDA_ARCHS=90a,100f \ # -> dist/liger_cute_kernels--py3-none-manylinux_2_35_x86_64.whl ``` +For a Hopper-only CUDA 13 build, select that toolkit's `nvcc`, install +`nvidia-nvshmem-cu13==3.6.5` in a separate environment, and use +`LIGER_CUTE_CUDA_ARCHS=90a` instead. +NVSHMEM must be installed **before** compiling from source; pip's extras are +runtime dependencies and do not provision native build prerequisites before +the wheel-build step. After provisioning the prerequisites, a source install +can use `pip install --no-build-isolation ".[cu12]"` or `".[cu13]"`. + For a native NVSHMEM install: ```bash @@ -617,4 +654,5 @@ The core should export only `liger_cute_*` symbols and have no direct `tvm_ffi.load_module`. - NVSHMEM is not copied into the LCK wheel. The loader preloads `libnvshmem_host.so.3` from the separately installed - `nvidia-nvshmem-cu12==3.6.5` package. + `nvidia-nvshmem-cu12==3.6.5` or `nvidia-nvshmem-cu13==3.6.5` package, + matching the native core's build toolkit. diff --git a/liger_cute_kernels/csrc/core/CMakeLists.txt b/liger_cute_kernels/csrc/core/CMakeLists.txt index 72224338..d237f3f6 100644 --- a/liger_cute_kernels/csrc/core/CMakeLists.txt +++ b/liger_cute_kernels/csrc/core/CMakeLists.txt @@ -104,9 +104,8 @@ if(LIGER_CUTE_ENABLE_SM90_NONRDC_MOE) "-I${CMAKE_CURRENT_SOURCE_DIR}/include" "-I${CMAKE_CURRENT_SOURCE_DIR}/src/moe" "-I${CMAKE_CURRENT_SOURCE_DIR}/src/fused_scaled_linear_cross_entropy" - "-I${NVSHMEM_INCLUDE_DIR}" - "-I${CUDAToolkit_INCLUDE_DIRS}") - foreach(_include_dir IN LISTS CUTLASS_INCLUDE_DIRS) + "-I${NVSHMEM_INCLUDE_DIR}") + foreach(_include_dir IN LISTS CUDAToolkit_INCLUDE_DIRS CUTLASS_INCLUDE_DIRS) list(APPEND _nonrdc_moe_include_args "-I${_include_dir}") endforeach() file(GLOB_RECURSE _nonrdc_moe_dependencies diff --git a/liger_cute_kernels/liger_cute_kernels/__init__.py b/liger_cute_kernels/liger_cute_kernels/__init__.py index d2a576cd..fc68d125 100644 --- a/liger_cute_kernels/liger_cute_kernels/__init__.py +++ b/liger_cute_kernels/liger_cute_kernels/__init__.py @@ -11,8 +11,8 @@ extension and its support libraries side by side:: liger_moe_sm90_nonrdc.cubin # optional local/IB Hopper module The Python API loads the core through TVM FFI, so the runtime boundary is the -torch-free core ABI rather than a Torch extension. NVSHMEM is installed through -the separate ``nvidia-nvshmem-cu12`` runtime dependency. +torch-free core ABI rather than a Torch extension. NVSHMEM is supplied separately; +the ``cu12`` and ``cu13`` extras install the matching pinned runtime dependency. The package currently includes expert-parallel MoE and tensor-parallel fused scaled linear cross-entropy kernels. Consumers should go through diff --git a/liger_cute_kernels/liger_cute_kernels/tvm_ffi.py b/liger_cute_kernels/liger_cute_kernels/tvm_ffi.py index ca3ca4e1..488c881d 100644 --- a/liger_cute_kernels/liger_cute_kernels/tvm_ffi.py +++ b/liger_cute_kernels/liger_cute_kernels/tvm_ffi.py @@ -92,7 +92,10 @@ def _load_nvshmem_libraries(pkg_dir: Path) -> None: None, ) if host is None: - raise ImportError("NVSHMEM runtime not found; install nvidia-nvshmem-cu12==3.6.5") + raise ImportError( + "NVSHMEM runtime not found; install 'liger-cute-kernels[cu12]' or " + "'liger-cute-kernels[cu13]' to match the CUDA version used to build the native wheel" + ) ctypes.CDLL(str(host), mode=ctypes.RTLD_GLOBAL) uid_bootstrap = host.parent / "nvshmem_bootstrap_uid.so.3" if uid_bootstrap.is_file(): diff --git a/liger_cute_kernels/setup.py b/liger_cute_kernels/setup.py index cf41462a..1b7b7d55 100644 --- a/liger_cute_kernels/setup.py +++ b/liger_cute_kernels/setup.py @@ -8,8 +8,8 @@ into ``liger_kernel``. The native wheel: * builds torch-free SM90a and SM100f cores with TVM FFI exports, * selects the matching core at runtime from the active GPU capability, - * installs NVSHMEM as a separate CUDA 12 dependency instead of copying its - shared libraries into this wheel, and + * offers CUDA 12/13 NVSHMEM dependencies through optional extras instead of + copying their shared libraries into this wheel, and * uses the same public version as the top-level ``liger-kernel`` package. ``liger_kernel.ops.cute`` (from the liger wheel) imports ``liger_cute_kernels.tvm_ffi`` @@ -51,8 +51,11 @@ setup( install_requires=[ "torch", "apache-tvm-ffi", - "nvidia-nvshmem-cu12==3.6.5", ], + extras_require={ + "cu12": ["nvidia-nvshmem-cu12==3.6.5"], + "cu13": ["nvidia-nvshmem-cu13==3.6.5"], + }, # Self-contained package: its __init__.py is packaged by build_py and the # .so are placed beside it by LckBuildExt. The marker extension makes this a # platform wheel, while LckBdistWheel records that the TVM FFI core is not diff --git a/liger_cute_kernels/test/test_cute_build.py b/liger_cute_kernels/test/test_cute_build.py index 90b930b9..9890d6eb 100644 --- a/liger_cute_kernels/test/test_cute_build.py +++ b/liger_cute_kernels/test/test_cute_build.py @@ -1,10 +1,14 @@ from __future__ import annotations import importlib.util +import runpy +import shutil +import subprocess from pathlib import Path import pytest +import setuptools _MODULE_PATH = Path(__file__).resolve().parents[1] / "cute_build.py" _SPEC = importlib.util.spec_from_file_location("lck_cute_build", _MODULE_PATH) @@ -13,6 +17,21 @@ cute_build = importlib.util.module_from_spec(_SPEC) _SPEC.loader.exec_module(cute_build) +def test_wheel_offers_optional_cuda_nvshmem_dependencies(monkeypatch): + metadata = {} + monkeypatch.setattr(setuptools, "setup", lambda **kwargs: metadata.update(kwargs)) + monkeypatch.syspath_prepend(str(_MODULE_PATH.parent)) + + runpy.run_path(str(_MODULE_PATH.parent / "setup.py"), run_name="__main__") + + nvshmem = [requirement for requirement in metadata["install_requires"] if requirement.startswith("nvidia-nvshmem-")] + assert nvshmem == [] + assert metadata["extras_require"] == { + "cu12": ["nvidia-nvshmem-cu12==3.6.5"], + "cu13": ["nvidia-nvshmem-cu13==3.6.5"], + } + + def test_prepare_nvshmem_home_adapts_versioned_pypi_layout(tmp_path, monkeypatch): home = tmp_path / "site-packages" / "nvidia" / "nvshmem" (home / "include").mkdir(parents=True) @@ -130,6 +149,33 @@ def test_cmake_args_disable_sm90_nonrdc_moe(monkeypatch): assert "-DLIGER_CUTE_ENABLE_SM90_NONRDC_MOE=OFF" in cute_build._cmake_base_args() +@pytest.mark.parametrize("cuda_includes", ["/cuda/include", "/cuda/include;/cuda/include/cccl"]) +def test_nonrdc_cuda_include_list(tmp_path, cuda_includes): + cmake = shutil.which("cmake") + if cmake is None: + pytest.skip("cmake is required to exercise CUDA include-list expansion") + source = (_MODULE_PATH.parent / "csrc/core/CMakeLists.txt").read_text() + start = source.index(" set(_nonrdc_moe_include_args") + end = source.index(" endforeach()", start) + len(" endforeach()") + script = tmp_path / "includes.cmake" + output = tmp_path / "includes.txt" + script.write_text( + f'set(CUDAToolkit_INCLUDE_DIRS "{cuda_includes}")\n' + 'set(CUTLASS_INCLUDE_DIRS "/cutlass/include;/cutlass/tools/util/include")\n' + + source[start:end] + + f'\nfile(WRITE "{output.as_posix()}" "${{_nonrdc_moe_include_args}}")\n' + ) + + subprocess.run([cmake, "-P", str(script)], check=True, capture_output=True, text=True) + + arguments = output.read_text().split(";") + assert all(argument.startswith("-I") for argument in arguments) + for directory in cuda_includes.split(";"): + assert f"-I{directory}" in arguments + assert "-I/cutlass/include" in arguments + assert "-I/cutlass/tools/util/include" in arguments + + def test_stage_optional_nonrdc_cubin(tmp_path): source = tmp_path / "build" / "csrc" / "core" output = tmp_path / "output" diff --git a/liger_cute_kernels/test/test_tvm_ffi_loader.py b/liger_cute_kernels/test/test_tvm_ffi_loader.py index ec06aec3..04894ff7 100644 --- a/liger_cute_kernels/test/test_tvm_ffi_loader.py +++ b/liger_cute_kernels/test/test_tvm_ffi_loader.py @@ -64,3 +64,11 @@ def test_nvshmem_loader_uses_dependency_package(tmp_path, monkeypatch): tvm_ffi._load_nvshmem_libraries(tmp_path / "wheel-package") assert [path for path, _ in loaded] == [str(host), str(uid)] + + +def test_nvshmem_loader_missing_runtime_recommends_cuda_extras(tmp_path, monkeypatch): + monkeypatch.setattr(tvm_ffi, "_nvshmem_library_dirs", lambda _: [tmp_path]) + monkeypatch.setattr(tvm_ffi, "_NVSHMEM_LIBS_LOADED", False) + + with pytest.raises(ImportError, match=r"liger-cute-kernels\[cu12\].*liger-cute-kernels\[cu13\].*build"): + tvm_ffi._load_nvshmem_libraries(tmp_path)