Fix candidate wheel assembly

.github/release/carriers.py

@@ -30,6 +30,7 @@
     "transformers-kt": "transformers_kt_sgl_kernel_payload",
     "sglang-kt": "sglang_kt_sgl_kernel_payload",
     "ktransformers": "ktransformers_sgl_kernel_payload",
+    "accelerate-kt": "accelerate_kt_sgl_kernel_payload",
 }
 LARGE = ("flash_ops.abi3.so", "sm100/common_ops.abi3.so")
 
@@ -44,7 +45,12 @@ def unpack(wheel, root):
     with zipfile.ZipFile(wheel) as archive:
         names = archive.namelist()
         require(len(names) == len(set(names)), "Duplicate ZIP entries")
-        records = [name for name in names if name.endswith(".dist-info/RECORD")]
+        records = [
+            name
+            for name in names
+            if name.endswith(".dist-info/RECORD")
+            and len(PurePosixPath(name).parts) == 2
+        ]
         require(len(records) == 1, "Require one wheel RECORD")
         rows = list(csv.reader(io.StringIO(archive.read(records[0]).decode())))
         record = {row[0]: row[1:] for row in rows}
@@ -139,15 +145,18 @@ def retag(root, python, abi):
 
 def archive_payload(sgl, output):
     hashes = {}
+    paths = [sgl / "sgl_kernel" / name for name in LARGE]
+    for directory in sorted(sgl.glob("*.libs")):
+        paths.extend(path for path in sorted(directory.rglob("*")) if path.is_file())
     with (
         output.open("wb") as raw,
         gzip.GzipFile(
             filename="", mode="wb", fileobj=raw, mtime=0, compresslevel=9
         ) as compressed,
     ):
         with tarfile.open(fileobj=compressed, mode="w|") as archive:
-            for name in LARGE:
-                source = sgl / "sgl_kernel" / name
+            for source in paths:
+                name = source.relative_to(sgl).as_posix()
                 hashes[name] = sha256(source)
                 info = archive.gettarinfo(str(source), arcname=name)
                 info.uid = info.gid = info.mtime = 0
@@ -189,7 +198,7 @@ def binary_evidence(roots):
     )
     kt = [entry for name, entry in evidence.items() if name.startswith("kt-kernel/")]
     require(
-        any(required <= normalize(entry["sass"]) for entry in kt),
+        not kt or any(required <= normalize(entry["sass"]) for entry in kt),
         "KT CUDA extension is missing required SASS architectures",
     )
     return evidence
@@ -201,7 +210,7 @@ def assemble(raw, output, evidence_dir):
         inspect_wheel(path) | {"path": path} for path in sorted(raw.glob("*.whl"))
     ]
     by_name = {entry["name"]: entry for entry in entries}
-    expected = set(MODULES) | {"accelerate-kt", "sgl-kernel-kt"}
+    expected = set(MODULES) | {"sgl-kernel-kt"}
     require(
         set(by_name) == expected and len(entries) == len(expected),
         "Need six fresh inputs",
@@ -224,24 +233,6 @@ def assemble(raw, output, evidence_dir):
             unpack(entry["path"], roots[name])
         save_json(evidence_dir / "cuda-binaries.json", binary_evidence(roots))
         sgl = roots["sgl-kernel-kt"]
-        # Lazy objects are extracted into a cache, not site-packages. An
-        # auditwheel-renamed dependency resolved via $ORIGIN would break there.
-        # Torch/CUDA-runtime libraries are supplied by the pinned torch wheel;
-        # reject any other bundled dependency rather than emitting a broken wheel.
-        bundled = {
-            path.name
-            for path in sgl.rglob("*.so*")
-            if any(part.endswith(".libs") for part in path.parts)
-        }
-        for name in LARGE:
-            dynamic = subprocess.check_output(
-                ["readelf", "--dynamic", str(sgl / "sgl_kernel" / name)], text=True
-            )
-            needed = set(re.findall(r"Shared library: \[([^]]+)\]", dynamic))
-            require(
-                not (needed & bundled),
-                "Lazy payload depends on a relocated auditwheel library; fix the main build linkage",
-            )
         for name in ("payload_runtime.py", "load_utils.py", "flash_attn.py"):
             require(
                 (sgl / "sgl_kernel" / name).is_file(),
@@ -261,6 +252,7 @@ def assemble(raw, output, evidence_dir):
             f"VERSION = {by_name['sgl-kernel-kt']['version']!r}\n"
             f"ARCHIVE_SHA256 = {sha256(archive)!r}\n"
             f"PAYLOAD_MODULES = {tuple(MODULES.values())!r}\nFILES = {hashes!r}\n"
+            f"BINARIES = {dict((name, 'sgl_kernel/' + name) for name in LARGE)!r}\n"
         )
         kt = roots["kt-kernel"]
         kt_dist = next(kt.glob("*.dist-info"))
@@ -355,12 +347,6 @@ def assemble(raw, output, evidence_dir):
                     "Carrier changed dependency metadata",
                 )
             require(remaining == 0 and not source.read(1), "Incomplete payload split")
-        accelerate = by_name["accelerate-kt"]["path"]
-        require(
-            accelerate.stat().st_size < LIMIT,
-            "Accelerate wheel exceeds PyPI size limit",
-        )
-        shutil.copyfile(accelerate, output / accelerate.name)
         save_json(
             evidence_dir / "assembly.json",
             {

.github/release/four_main.py

@@ -205,7 +205,9 @@ def inspect_wheel(path: Path) -> dict:
     name, version, _, tags = parse_wheel_filename(path.name)
     with zipfile.ZipFile(path) as wheel:
         metadata_files = [
-            p for p in wheel.namelist() if p.endswith(".dist-info/METADATA")
+            p
+            for p in wheel.namelist()
+            if p.endswith(".dist-info/METADATA") and len(PurePosixPath(p).parts) == 2
         ]
         if len(metadata_files) != 1:
             raise ValueError(f"Expected one METADATA in {path.name}")

.github/release/test_carriers.py

@@ -1,8 +1,10 @@
 """Exercise carrier assembly with tiny synthetic wheels, not CUDA execution."""
 
 import sys
+import tarfile
 import zipfile
 from pathlib import Path
+from types import SimpleNamespace
 
 import pytest
 
@@ -14,7 +16,7 @@
 def raw_wheels(tmp_path):
     raw = tmp_path / "raw"
     raw.mkdir()
-    for package in (*carriers.MODULES, "accelerate-kt", "sgl-kernel-kt"):
+    for package in (*carriers.MODULES, "sgl-kernel-kt"):
         root = tmp_path / package
         root.mkdir()
         dist = root / (package.replace("-", "_") + "-1.0.dist-info")
@@ -73,6 +75,8 @@ def test_fresh_carriers_preserve_runtime_versions_and_sm90(tmp_path, monkeypatch
         assert "sgl_kernel/_payload_manifest.py" in wheel.namelist()
     with zipfile.ZipFile(next(output.glob("ktransformers-*.whl"))) as wheel:
         assert "sgl_kernel/sm90/common_ops.abi3.so" in wheel.namelist()
+    with zipfile.ZipFile(next(output.glob("accelerate_kt-*.whl"))) as wheel:
+        assert "accelerate_kt_sgl_kernel_payload/payload.part" in wheel.namelist()
     # Every final RECORD is independently checked, including regenerated payloads.
     for index, path in enumerate(output.iterdir()):
         carriers.unpack(path, tmp_path / f"verify-{index}")
@@ -96,3 +100,65 @@ def test_size_limit_fails_without_dropping_architectures(tmp_path, monkeypatch):
     evidence.mkdir()
     with pytest.raises(ValueError, match="capacity"):
         carriers.assemble(raw, tmp_path / "final", evidence)
+
+
+def test_payload_preserves_wheel_relative_libraries(tmp_path):
+    raw_wheels(tmp_path)
+    root = tmp_path / "sgl-kernel-kt"
+    library = root / "sgl_kernel_kt.libs/libnuma.so.1"
+    library.parent.mkdir()
+    library.write_bytes(b"dependency")
+    target = tmp_path / "payload.tar.gz"
+    hashes = carriers.archive_payload(root, target)
+    assert set(hashes) == {"sgl_kernel/" + name for name in carriers.LARGE} | {
+        "sgl_kernel_kt.libs/libnuma.so.1"
+    }
+    with tarfile.open(target) as archive:
+        assert set(archive.getnames()) == set(hashes)
+
+
+def test_embedded_metadata_does_not_shadow_wheel_metadata(tmp_path):
+    raw_wheels(tmp_path)
+    root = tmp_path / "accelerate-kt"
+    nested = root / "accelerate_kt/vendor/example.dist-info"
+    nested.mkdir(parents=True)
+    for name in ("METADATA", "RECORD"):
+        (nested / name).write_text("vendored metadata")
+    wheel = tmp_path / "accelerate_kt-1.0-py3-none-any.whl"
+    carriers.pack(root, wheel)
+    assert inspect_wheel(wheel)["name"] == "accelerate-kt"
+    carriers.unpack(wheel, tmp_path / "unpacked")
+
+
+@pytest.mark.parametrize("kt_cuda", [False, True])
+def test_cpu_only_kt_is_accepted_without_weakening_cuda_audit(
+    tmp_path, monkeypatch, kt_cuda
+):
+    kt, sgl = tmp_path / "kt", tmp_path / "sgl"
+    kt.mkdir()
+    (sgl / "sgl_kernel/sm100").mkdir(parents=True)
+    (kt / "extension.so").touch()
+    (sgl / "sgl_kernel/sm100/common_ops.abi3.so").touch()
+    monkeypatch.setattr(
+        carriers.subprocess,
+        "check_output",
+        lambda args, **kw: (
+            ".nv_fatbin" if kt_cuda or sgl in Path(args[-1]).parents else ""
+        ),
+    )
+    monkeypatch.setattr(
+        carriers.subprocess,
+        "run",
+        lambda args, **kw: SimpleNamespace(
+            stdout=(
+                "sm_80 sm_86 sm_89 sm_90 sm_120"
+                if sgl in Path(args[-1]).parents
+                else "sm_80"
+            )
+        ),
+    )
+    if kt_cuda:
+        with pytest.raises(ValueError, match="KT CUDA extension"):
+            carriers.binary_evidence({"kt-kernel": kt, "sgl-kernel-kt": sgl})
+    else:
+        carriers.binary_evidence({"kt-kernel": kt, "sgl-kernel-kt": sgl})

.github/workflows/release-four-main.yml

@@ -229,6 +229,10 @@ jobs:
                       '--exclude', 'libtorch_cuda.so', '--exclude', 'libtorch_python.so',
                       '--exclude', 'libc10.so', '--exclude', 'libc10_cuda.so',
                       '--exclude', 'libcudart.so.12',
+                      '--exclude', 'libcublas.so.12',
+                      '--exclude', 'libcublasLt.so.12',
+                      '--exclude', 'libcurand.so.10',
+                      '--exclude', 'libnvrtc.so.12',
                       '-w', str(work / 'repaired'),
                   ], check=True)
               else: