Skip to content

[TIRx][CUDA] Allow newer CUTLASS packages for IKET #20164

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jinhongyii wants to merge 1 commit into
base: main
Choose a base branch
from
+87 −102
Open

[TIRx][CUDA] Allow newer CUTLASS packages for IKET #20164

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
114 changes: 40 additions & 74 deletions python/tvm/backend/cuda/iket.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,6 @@
from __future__ import annotations

import functools
import hashlib
import importlib.util
import inspect
import json
Expand Down Expand Up @@ -59,48 +58,19 @@
_OUTPUT_TAIL_LINES = 100
_TERMINATION_GRACE_SECONDS = 5.0

# Hashes are SHA-256 digests of files in the public 4.6.0 wheels. Only
# ABI-independent runtime/compiler binaries are pinned so this profile works
# with every Python version supported by that CUTLASS DSL release.
_OFFICIAL_PROFILES = {
"cutlass-4.6.0": {
"nvrtc_version": (13, 2),
"versions": {
"minimum_versions": {
"nvidia-cutlass-dsl": "4.6.0",
"nvidia-cutlass-dsl-libs-base": "4.6.0",
"nvidia-cutlass-dsl-libs-core": "4.6.0",
"nvidia-cutlass-dsl-libs-cu13": "4.6.0",
},
"exact_versions": {
"nvidia-cuda-nvdisasm": "13.3.73",
"nvidia-cuda-nvrtc": "13.2.78",
},
"files": {
"nvidia-cutlass-dsl-libs-base": {
"nvidia_cutlass_dsl/dsl_packages/iket/libiket_cubin_info.so": (
"7ee839130c6bd129b04908a807c066118a459ebea644a59ecb6e41fbb323c103"
),
"nvidia_cutlass_dsl/dsl_packages/iket/profiler/libsmodel_injection.so": (
"83be54bd06e2cd82b2f6c17bbee6c925d049acae8d880242d4a5d5509a29e122"
),
},
"nvidia-cutlass-dsl-libs-cu13": {
"nvidia_cutlass_dsl/cu13/lib/libcute_dsl_runtime.so": (
"2fa9809047485ae420ca99cab0678846de692e9608a179b0020834994311dd2f"
),
},
"nvidia-cuda-nvdisasm": {
"nvidia/cu13/bin/nvdisasm": (
"5842e6adf9e232c9503a804915f158a576473e542577c070da3be49390474140"
),
},
"nvidia-cuda-nvrtc": {
"nvidia/cu13/lib/libnvrtc.so.13": (
"c673cf3b5099d83b98a388a2bb21e5d6f481be3c4bb956e2d74c39cb714d8c63"
),
"nvidia/cu13/lib/libnvrtc-builtins.so.13.2": (
"6b1c571cc730d5fcfd57f322e1fa7e0e65de7454b2239ff6d552a09b82d47dbe"
),
},
},
}
}

Expand Down Expand Up @@ -201,12 +171,17 @@ def _profile_error(message: str) -> IketProfileError:
)


def _sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as file_obj:
for chunk in iter(lambda: file_obj.read(1 << 20), b""):
digest.update(chunk)
return digest.hexdigest()
def _is_newer_release(actual: str, expected: str) -> bool:
"""Compare the numeric release versions published by NVIDIA wheels."""
try:
actual_parts = tuple(int(part) for part in actual.split("."))
expected_parts = tuple(int(part) for part in expected.split("."))
except ValueError:
return False
width = max(len(actual_parts), len(expected_parts))
return actual_parts + (0,) * (width - len(actual_parts)) > expected_parts + (0,) * (
width - len(expected_parts)
)


def _validate_run_iket_entrypoint() -> str:
Expand All @@ -220,7 +195,9 @@ def _validate_run_iket_entrypoint() -> str:
and item.value == "iket.cli.main:entrypoint"
for item in entry_points
):
raise _profile_error("the run-iket entry point does not match the locked CUTLASS profile")
raise _profile_error(
"the run-iket entry point does not match the supported CUTLASS profile"
)
return executable


Expand All @@ -240,48 +217,41 @@ def _validate_nvrtc_version(expected_version: tuple[int, int]) -> None:


def _validate_official_installation(profile_name: str) -> str:
"""Validate pinned host-side packages and return the official executable."""
"""Validate host-side package versions and return the official executable."""
if profile_name not in _OFFICIAL_PROFILES:
raise _profile_error(f"unsupported profile {profile_name!r}; expected {_DEFAULT_PROFILE!r}")
profile_config = _OFFICIAL_PROFILES[profile_name]
distributions = {}
for distribution_name, expected_version in profile_config["versions"].items():
try:
distribution = metadata.distribution(distribution_name)
except metadata.PackageNotFoundError as err:
raise _profile_error(
f"{distribution_name}=={expected_version} is not installed"
) from err
if distribution.version != expected_version:
raise _profile_error(
f"{distribution_name} must be {expected_version}, got {distribution.version}"
)
distributions[distribution_name] = distribution

for distribution_name, expected_files in profile_config["files"].items():
distribution = distributions[distribution_name]
for relative_path, expected_digest in expected_files.items():
path = Path(distribution.locate_file(relative_path))
if not path.is_file():
raise _profile_error(f"profile binary is missing: {relative_path}")
actual_digest = _sha256(path)
if actual_digest != expected_digest:
version_groups = (
(profile_config["minimum_versions"], True),
(profile_config["exact_versions"], False),
)
for versions, allow_newer in version_groups:
for distribution_name, expected_version in versions.items():
try:
distribution = metadata.distribution(distribution_name)
except metadata.PackageNotFoundError as err:
operator = ">=" if allow_newer else "=="
raise _profile_error(
f"profile binary hash mismatch for {relative_path}: {actual_digest}"
f"{distribution_name}{operator}{expected_version} is not installed"
) from err
if distribution.version != expected_version and not (
allow_newer and _is_newer_release(distribution.version, expected_version)
):
requirement = f"{expected_version} or newer" if allow_newer else expected_version
raise _profile_error(
f"{distribution_name} must be {requirement}, got {distribution.version}"
)

executable = _validate_run_iket_entrypoint()
_validate_nvrtc_version(profile_config["nvrtc_version"])
return executable


def _validate_injection_environment(expected_injection_digest: str) -> None:
def _validate_injection_environment() -> None:
injection_value = os.environ.get("CUDA_INJECTION64_PATH")
injection_path = Path(injection_value) if injection_value else None
if injection_path is None or not injection_path.is_file():
raise _profile_error("CUDA_INJECTION64_PATH was not supplied by run-iket")
if _sha256(injection_path) != expected_injection_digest:
raise _profile_error("CUDA_INJECTION64_PATH does not match the locked run-iket binary")

config_value = os.environ.get("SMODEL_INJECTION_CONFIG")
config_path = Path(config_value) if config_value else None
Expand Down Expand Up @@ -314,16 +284,12 @@ def _validate_official_environment() -> str:
f"{_PROFILE_ENV} must be set to {_DEFAULT_PROFILE}, got {profile_name!r}"
)
executable = _validate_official_installation(profile_name)
profile_config = _OFFICIAL_PROFILES[profile_name]
injection_digest = profile_config["files"]["nvidia-cutlass-dsl-libs-base"][
"nvidia_cutlass_dsl/dsl_packages/iket/profiler/libsmodel_injection.so"
]
_validate_injection_environment(injection_digest)
_validate_injection_environment()
return executable


def validate_official_environment() -> None:
"""Validate the exact official runtime before an instrumented CUBIN is loaded."""
"""Validate the supported official runtime before an instrumented CUBIN is loaded."""
_validate_official_environment()


Expand Down Expand Up @@ -469,7 +435,7 @@ def _child_environment(
child_env[_PROFILE_ENV] = profile_name
# LowerIket also requires the two run-iket injection variables before it
# honors this marker. This enables ordinary TIRx JIT compilation only in
# children started by this locked profiling entry point.
# children started by this validated profiling entry point.
child_env[_INJECTED_CHILD_ENABLE_ENV] = "1"
return child_env

Expand Down
75 changes: 47 additions & 28 deletions tests/python/tirx/iket/test_iket_profiler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for IKET lowering, metadata, installation locks, and trace contracts."""
"""Tests for IKET lowering, metadata, installation versions, and trace contracts."""

import hashlib
import importlib
Expand Down Expand Up @@ -769,35 +769,22 @@ def test_proxy_fails_closed_and_forbids_export(tmp_path, monkeypatch):
assert executable._executable._jitted_mod is None # pylint: disable=protected-access


def test_environment_validation_is_not_process_cached(tmp_path, monkeypatch):
def test_environment_validation_is_not_process_cached(monkeypatch):
from tvm.tirx.cuda import iket as _iket_official

injection_path = tmp_path / "libsmodel_injection.so"
injection_path.write_bytes(b"locked")
injection_relative = "nvidia_cutlass_dsl/dsl_packages/iket/profiler/libsmodel_injection.so"
profile = {
"nvrtc_version": (13, 2),
"versions": {"nvidia-cutlass-dsl-libs-base": "4.6.0"},
"files": {
"nvidia-cutlass-dsl-libs-base": {
injection_relative: hashlib.sha256(b"locked").hexdigest()
}
},
"minimum_versions": {"nvidia-cutlass-dsl-libs-base": "4.6.0"},
"exact_versions": {},
}

class FakeDistribution:
version = "4.6.0"

@staticmethod
def locate_file(_relative_path):
return injection_path

monkeypatch.setitem(_iket_official._OFFICIAL_PROFILES, "cutlass-4.6.0", profile)
monkeypatch.setattr(_iket_official.metadata, "distribution", lambda _name: FakeDistribution())
monkeypatch.setattr(_iket_official, "_validate_run_iket_entrypoint", lambda: None)
monkeypatch.setattr(
_iket_official, "_validate_injection_environment", lambda _expected_digest: None
)
monkeypatch.setattr(_iket_official, "_validate_injection_environment", lambda: None)
monkeypatch.setattr(_iket_official, "_validate_nvrtc_version", lambda _version: None)
monkeypatch.setenv("TVM_IKET_OFFICIAL_PROFILE", "cutlass-4.6.0")

Expand All @@ -807,18 +794,55 @@ def locate_file(_relative_path):
_iket_official.validate_official_environment()


@pytest.mark.parametrize(
("version", "expected_error"),
(
("4.5.0", "must be 4.6.0 or newer"),
("4.6.0", None),
("4.6.2", None),
),
)
def test_official_installation_accepts_newer_cutlass(monkeypatch, version, expected_error):
from tvm.tirx.cuda import iket as _iket_official

profile = {
"nvrtc_version": (13, 2),
"minimum_versions": {"nvidia-cutlass-dsl-libs-base": "4.6.0"},
"exact_versions": {},
}

class FakeDistribution:
pass

distribution = FakeDistribution()
distribution.version = version
monkeypatch.setitem(_iket_official._OFFICIAL_PROFILES, "cutlass-4.6.0", profile)
monkeypatch.setattr(_iket_official.metadata, "distribution", lambda _name: distribution)
monkeypatch.setattr(_iket_official, "_validate_run_iket_entrypoint", lambda: None)
monkeypatch.setattr(_iket_official, "_validate_nvrtc_version", lambda _version: None)

if expected_error:
with pytest.raises(RuntimeError, match=expected_error):
_iket_official._validate_official_installation( # pylint: disable=protected-access
"cutlass-4.6.0"
)
else:
_iket_official._validate_official_installation( # pylint: disable=protected-access
"cutlass-4.6.0"
)


def test_injection_environment_accepts_run_iket_two_passes(tmp_path, monkeypatch):
from tvm.tirx.cuda import iket as _iket_official

injection = tmp_path / "libsmodel_injection.so"
injection.write_bytes(b"locked-injection")
expected_digest = hashlib.sha256(injection.read_bytes()).hexdigest()
injection.write_bytes(b"injection")
config_path = tmp_path / "config.json"
monkeypatch.setenv("CUDA_INJECTION64_PATH", str(injection))
monkeypatch.setenv("SMODEL_INJECTION_CONFIG", str(config_path))

config_path.write_text(json.dumps({"toolName": "tracker", "toolConfig": {}}))
_iket_official._validate_injection_environment(expected_digest) # pylint: disable=protected-access
_iket_official._validate_injection_environment() # pylint: disable=protected-access

instrument = tmp_path / "instrument.config.json"
instrument.write_text("{}", encoding="utf-8")
Expand All @@ -831,15 +855,10 @@ def test_injection_environment_accepts_run_iket_two_passes(tmp_path, monkeypatch
),
encoding="utf-8",
)
_iket_official._validate_injection_environment(expected_digest) # pylint: disable=protected-access

with pytest.raises(RuntimeError, match="locked run-iket binary"):
_iket_official._validate_injection_environment("0" * 64) # pylint: disable=protected-access
_iket_official._validate_injection_environment() # pylint: disable=protected-access
config_path.write_text(json.dumps({"toolName": "other"}), encoding="utf-8")
with pytest.raises(RuntimeError, match="not generated by run-iket profile"):
_iket_official._validate_injection_environment( # pylint: disable=protected-access
expected_digest
)
_iket_official._validate_injection_environment() # pylint: disable=protected-access


def test_cutlass_4_6_0_oracle_manifest_integrity():
Expand Down
Loading