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-<version>-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 <[email protected]> Co-authored-by: Copilot <[email protected]>
2026-10-12 02:17:47 +00:00 · 2026-09-25 16:27:26 +00:00
parent 40a9d8a61b
commit be2a4eb14c
9 changed files with 139 additions and 23 deletions
@@ -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

@@ -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
@@ -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-<version>-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-<liger-version>-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.
@@ -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
@@ -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
@@ -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():
@@ -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
@@ -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"
@@ -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)