ROCm bundles: ship rocBLAS's split kernel packs inside the bundle

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

.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

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
+`<rocm>/.kpack/blas_lib_<gfx>.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())

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()