From e9af246d3e78c1697451935b2a44ac2cd5336aea Mon Sep 17 00:00:00 2001 From: Hongyi Jin Date: Fri, 21 Aug 2026 22:00:29 -0400 Subject: [PATCH] [TIRx][CUDA] Allow newer CUTLASS packages for IKET The official IKET profile currently pins CUTLASS 4.6.0 package versions and binary hashes, which rejects newer CUTLASS installations even when they retain the supported interface.\n\nTreat the four CUTLASS packages as requiring version 4.6.0 or newer while preserving the existing exact NVRTC and NVDisasm requirements. Remove binary hash pinning, and cover the minimum-version boundary in the installation validation tests. --- python/tvm/backend/cuda/iket.py | 114 +++++++------------ tests/python/tirx/iket/test_iket_profiler.py | 75 +++++++----- 2 files changed, 87 insertions(+), 102 deletions(-) diff --git a/python/tvm/backend/cuda/iket.py b/python/tvm/backend/cuda/iket.py index c1edbed2effb..62d65c4cce72 100644 --- a/python/tvm/backend/cuda/iket.py +++ b/python/tvm/backend/cuda/iket.py @@ -24,7 +24,6 @@ from __future__ import annotations import functools -import hashlib import importlib.util import inspect import json @@ -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" - ), - }, - }, } } @@ -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: @@ -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 @@ -240,34 +217,29 @@ 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() @@ -275,13 +247,11 @@ def _validate_official_installation(profile_name: str) -> str: 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 @@ -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() @@ -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 diff --git a/tests/python/tirx/iket/test_iket_profiler.py b/tests/python/tirx/iket/test_iket_profiler.py index 4ad200e912b0..ba6f30f3d7f5 100644 --- a/tests/python/tirx/iket/test_iket_profiler.py +++ b/tests/python/tirx/iket/test_iket_profiler.py @@ -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 @@ -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") @@ -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") @@ -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():