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]>
This commit is contained in:
Jari Kolehmainen
2026-09-25 16:27:26 +00:00
committed by GitHub
co-authored by jkolehma Copilot
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