From f95b52162ed4b916a58d4f16b72662483cc55075 Mon Sep 17 00:00:00 2001 From: Daniel Han <23090290+danielhanchen@users.noreply.github.com> Date: Wed, 14 Oct 2026 23:21:21 -0700 Subject: [PATCH] ROCm bundles: ship rocBLAS's split kernel packs inside the bundle --- .github/workflows/unsloth-pr-set-lint.yml | 5 +- .github/workflows/unsloth-prebuilt-rocm.yml | 28 +++++++ scripts/unsloth/relocate_kpack.py | 83 +++++++++++++++++++++ scripts/unsloth/test_relocate_kpack.py | 76 +++++++++++++++++++ 4 files changed, 191 insertions(+), 1 deletion(-) create mode 100644 scripts/unsloth/relocate_kpack.py create mode 100644 scripts/unsloth/test_relocate_kpack.py diff --git a/.github/workflows/unsloth-pr-set-lint.yml b/.github/workflows/unsloth-pr-set-lint.yml index ca0aaf18714e..6c32ec0a485a 100644 --- a/.github/workflows/unsloth-pr-set-lint.yml +++ b/.github/workflows/unsloth-pr-set-lint.yml @@ -20,6 +20,7 @@ on: - scripts/unsloth/merge_checks.py - scripts/unsloth/carry_vintage.py - scripts/unsloth/sync_deletes.py + - scripts/unsloth/relocate_kpack.py - scripts/unsloth/test_*.py - scripts/unsloth/check_workflow_scalars.py - scripts/unsloth/check_workflow_triggers.py @@ -35,6 +36,7 @@ on: - scripts/unsloth/merge_checks.py - scripts/unsloth/carry_vintage.py - scripts/unsloth/sync_deletes.py + - scripts/unsloth/relocate_kpack.py - scripts/unsloth/test_*.py - scripts/unsloth/check_workflow_scalars.py - scripts/unsloth/check_workflow_triggers.py @@ -145,7 +147,8 @@ jobs: scripts/unsloth/test_carry_vintage.py \ scripts/unsloth/test_pin_contract.py \ scripts/unsloth/test_feature_matrix.py \ - scripts/unsloth/test_check_workflow_triggers.py; do + scripts/unsloth/test_check_workflow_triggers.py \ + scripts/unsloth/test_relocate_kpack.py; do echo "::group::$t" python3 "$t" || fail=1 echo "::endgroup::" diff --git a/.github/workflows/unsloth-prebuilt-rocm.yml b/.github/workflows/unsloth-prebuilt-rocm.yml index 145677342ffd..cf77bac41d71 100644 --- a/.github/workflows/unsloth-prebuilt-rocm.yml +++ b/.github/workflows/unsloth-prebuilt-rocm.yml @@ -384,6 +384,7 @@ jobs: set "mapped_target=%current_target%" ) echo Mapped target: %mapped_target% + echo MAPPED_TARGET=%mapped_target%>> "%GITHUB_ENV%" REM Set up environment variables and PATH set HIP_PATH=C:\opt\rocm @@ -512,6 +513,20 @@ jobs: exit 1 } + # The source tarball carries no scripts/unsloth (the upstream checkout drops it). + - name: Fetch kpack tooling + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0 + with: + path: kpack-tooling + sparse-checkout: scripts/unsloth + persist-credentials: false + + # rocBLAS / hipBLASLt keep their own kernels in ../.kpack, outside a flat bundle (unslothai/unsloth#13663). + - name: Ship the rocBLAS kernel packs + run: | + python kpack-tooling\scripts\unsloth\relocate_kpack.py --bundle llama.cpp\build\bin --kpack-src C:\opt\rocm\.kpack --archs "$env:MAPPED_TARGET" + if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE } + - name: Check the bundle carries every ROCm DLL it imports run: | python -m pip install --quiet --disable-pip-version-check pefile @@ -844,6 +859,7 @@ jobs: mapped_target="$current_target" fi echo "Mapped target: $mapped_target" + echo "MAPPED_TARGET=$mapped_target" >> "$GITHUB_ENV" # Create build directory cd llama.cpp @@ -1066,6 +1082,18 @@ jobs: [ -f "$file" ] && [ ! -L "$file" ] && patchelf --set-rpath '$ORIGIN' "$file" 2>/dev/null || true done + # The source tarball carries no scripts/unsloth (the upstream checkout drops it). + - name: Fetch kpack tooling + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0 + with: + path: kpack-tooling + sparse-checkout: scripts/unsloth + persist-credentials: false + + # After patchelf, which can leave a second copy of the reference; rewrites every copy (unslothai/unsloth#13663). + - name: Ship the rocBLAS kernel packs + run: python3 kpack-tooling/scripts/unsloth/relocate_kpack.py --bundle llama.cpp/build/bin --kpack-src /opt/rocm/.kpack --archs "$MAPPED_TARGET" + # Reject missing dependencies and runner libraries that mask an incomplete # bundle. Only glibc, libstdc++ and libgcc_s may resolve outside the bundle. - name: Check the bundle resolves its own dependencies diff --git a/scripts/unsloth/relocate_kpack.py b/scripts/unsloth/relocate_kpack.py new file mode 100644 index 000000000000..8045dd897a1a --- /dev/null +++ b/scripts/unsloth/relocate_kpack.py @@ -0,0 +1,83 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +"""Ship rocBLAS's split device code inside a flat ROCm bundle. + +TheRock strips the GPU kernels out of librocblas / libhipblaslt (`.hip_fatbin` is NOBITS) into +`/.kpack/blas_lib_.kpack`, and records where to find them in each library's +`.rocm_kpack_ref` as `../.kpack/blas_lib_@GFXARCH@.kpack`, relative to the library's directory. +Our bundles are flat, so that path points outside the bundle: every rocBLAS kernel that lives in +the library itself (not in rocblas/library) fails with hipErrorInvalidKernelFile, and llama.cpp +aborts in hipblasSgemm on the first F32 GEMM (unslothai/unsloth#13663). + +This copies the archives to `rocblas/k/` and rewrites the reference prefix `../.kpack/` to the +same-length `rocblas/k/`, so no msgpack length or section size changes. `rocblas/` already ships +next to the libraries and Studio copies that tree whole, so installs that predate this fix pick +the archives up too. Run it after the libraries are copied and before signing. +""" + +from __future__ import annotations + +import argparse +import re +import shutil +import sys +from pathlib import Path + +OLD_PREFIX = b"../.kpack/" +NEW_PREFIX = b"rocblas/k/" +ARCHIVE="blas_lib" +REF = re.compile(re.escape(OLD_PREFIX) + rb"([a-z_]+)_@GFXARCH@\.kpack") +assert len(OLD_PREFIX) == len(NEW_PREFIX) + + +def _libraries(bundle: Path): + for path in sorted(bundle.iterdir()): + if path.is_symlink() or not path.is_file(): + continue + if path.suffix.lower() == ".dll" or ".so" in path.name: + yield path + + +def relocate(bundle: Path, kpack_src: Path, archs: list[str]) -> list[str]: + old = OLD_PREFIX + ARCHIVE.encode() + b"_@GFXARCH@.kpack" + new = NEW_PREFIX + ARCHIVE.encode() + b"_@GFXARCH@.kpack" + referencing = [] + for lib in _libraries(bundle): + data = lib.read_bytes() + if any(m.group(1).decode() == ARCHIVE for m in REF.finditer(data)): + referencing.append((lib, data)) + if not referencing: + # A ROCm dist that keeps the kernels inside the libraries needs nothing shipped. + return [] + + dest = bundle / NEW_PREFIX.decode() + dest.mkdir(parents=True, exist_ok=True) + for arch in archs: + src = kpack_src / f"{ARCHIVE}_{arch}.kpack" + if not src.is_file(): + raise SystemExit(f"missing {src}: the ROCm dist has no {ARCHIVE} kernels for {arch}") + shutil.copy2(src, dest / src.name) + # Every copy: patchelf can leave a second one in a moved segment, and the loader reads by address. + for lib, data in referencing: + lib.write_bytes(data.replace(old, new)) + return [lib.name for lib, _ in referencing] + + +def main() -> int: + p = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + p.add_argument("--bundle", required=True, type=Path, help="flat bundle dir holding the ROCm libraries") + p.add_argument("--kpack-src", required=True, type=Path, help="the ROCm dist's .kpack dir") + p.add_argument("--archs", required=True, help="gfx targets, comma or semicolon separated") + a = p.parse_args() + archs = [x.strip() for x in re.split(r"[,;]", a.archs) if x.strip()] + patched = relocate(a.bundle, a.kpack_src, archs) + if patched: + print(f"kpack: {', '.join(archs)} -> {NEW_PREFIX.decode()}; references rewritten in {', '.join(patched)}") + else: + print(f"kpack: no library in {a.bundle} references ../.kpack/{ARCHIVE}; nothing to ship") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/unsloth/test_relocate_kpack.py b/scripts/unsloth/test_relocate_kpack.py new file mode 100644 index 000000000000..df38627619fc --- /dev/null +++ b/scripts/unsloth/test_relocate_kpack.py @@ -0,0 +1,76 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +"""Tests for relocate_kpack.py. Run: python3 scripts/unsloth/test_relocate_kpack.py + +Fake libraries carry the reference exactly as TheRock embeds it (msgpack map in .rocm_kpack_ref); +the rewrite must keep every byte but the 10-byte prefix, so the msgpack lengths stay valid. +""" +import subprocess +import sys +import tempfile +from pathlib import Path + +SCRIPT = Path(__file__).resolve().parent / "relocate_kpack.py" +FAILS = [] + +REF_BLAS = b'\x82\xb2kpack_search_paths\x91\xd9"../.kpack/blas_lib_@GFXARCH@.kpack\xabkernel_name\xd92math-libs/BLAS/rocBLAS/stage/lib/librocblas.so.5.8' +REF_SOLVER = b'\x82\xb2kpack_search_paths\x91\xd9$../.kpack/solver_lib_@GFXARCH@.kpack\xabkernel_name\xd9\x07solver!' + + +def check(name, cond, extra=""): + print(f"{'PASS' if cond else 'FAIL'} {name}" + (f" :: {extra}" if extra and not cond else "")) + if not cond: + FAILS.append(name) + + +def run(bundle, kpack, archs): + return subprocess.run([sys.executable, str(SCRIPT), "--bundle", str(bundle), "--kpack-src", str(kpack), "--archs", archs], + capture_output=True, text=True) + + +def fixture(libs): + d = Path(tempfile.mkdtemp(prefix="kp_")) + bundle, kpack = d / "bin", d / ".kpack" + bundle.mkdir() + kpack.mkdir() + for name, body in libs.items(): + (bundle / name).write_bytes(body) + for arch in ("gfx1100", "gfx1103", "gfx1151"): + (kpack / f"blas_lib_{arch}.kpack").write_bytes(b"KPACK" + arch.encode()) + return bundle, kpack + + +def main(): + head, tail = b"\x7fELF" + b"\0" * 60, b"\0" * 32 + # Two copies, as patchelf leaves them; the solver reference must stay untouched. + rocblas = head + REF_BLAS + b"\0" * 16 + REF_BLAS + tail + bundle, kpack = fixture({"librocblas.so.5.8": rocblas, "librocsolver.so.0": head + REF_SOLVER + tail, + "libggml-hip.so": head + tail, "rocblas.dll": rocblas}) + r = run(bundle, kpack, "gfx1100;gfx1103") + check("exit 0", r.returncode == 0, r.stderr) + out = (bundle / "librocblas.so.5.8").read_bytes() + check("same length", len(out) == len(rocblas)) + check("both copies rewritten", out.count(b'"rocblas/k/blas_lib_@GFXARCH@.kpack') == 2 and b"../.kpack/blas" not in out) + check("only the prefix changed", sum(a != b for a, b in zip(out, rocblas)) <= 2 * 10) + check("dll rewritten too", b"rocblas/k/blas_lib" in (bundle / "rocblas.dll").read_bytes()) + check("solver reference untouched", (bundle / "librocsolver.so.0").read_bytes() == head + REF_SOLVER + tail) + check("requested archs shipped", sorted(p.name for p in (bundle / "rocblas" / "k").iterdir()) + == ["blas_lib_gfx1100.kpack", "blas_lib_gfx1103.kpack"]) + + bundle, kpack = fixture({"librocblas.so": rocblas}) + r = run(bundle, kpack, "gfx1201") + check("missing arch archive fails", r.returncode != 0 and "gfx1201" in (r.stderr + r.stdout)) + + bundle, kpack = fixture({"librocblas.so": head + tail}) + r = run(bundle, kpack, "gfx1151") + check("unsplit dist is a no-op", r.returncode == 0 and not (bundle / "rocblas").exists(), r.stderr) + + if FAILS: + print(f"\n{len(FAILS)} failed: {FAILS}") + sys.exit(1) + print("\nall passed") + + +if __name__ == "__main__": + main()