This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new ad7afa6a12 [TIRx][CUDA] Allow newer CUTLASS packages for IKET (#20164)
ad7afa6a12 is described below
commit ad7afa6a129afd864ea7fba2405ed1abd1b29b3b
Author: Hongyi Jin <[email protected]>
AuthorDate: Mon Aug 24 00:44:28 2026 -0400
[TIRx][CUDA] Allow newer CUTLASS packages for IKET (#20164)
The official IKET profile currently requires the CUTLASS packages to be
exactly version 4.6.0 and locks their runtime files by hash. This
rejects newer package releases such as 4.6.2 even when they retain the
supported run-iket interface.
This change represents CUTLASS requirements as minimum versions and
accepts releases at or above 4.6.0. It preserves the existing exact
nvidia-cuda-nvrtc and nvidia-cuda-nvdisasm package requirements and the
NVRTC runtime check. Binary hash pinning is removed, while the run-iket
entry point and injected environment structure continue to be validated.
Testing:
- pre-commit run --files python/tvm/backend/cuda/iket.py
tests/python/tirx/iket/test_iket_profiler.py
- python -m pytest tests/python/tirx/iket/test_iket_profiler.py
tests/python/tirx/iket/test_iket_orchestration.py -q (71 passed, 2
skipped)
- Local installation validation with CUTLASS 4.6.2, NVDisasm 13.3.73,
and NVRTC 13.2.78/runtime 13.2
---
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 c1edbed2ef..62d65c4cce 100644
--- a/python/tvm/backend/cuda/iket.py
+++ b/python/tvm/backend/cuda/iket.py
@@ -24,7 +24,6 @@ capture processes execute the supplied callable and then exit.
from __future__ import annotations
import functools
-import hashlib
import importlib.util
import inspect
import json
@@ -59,48 +58,19 @@ _INJECTION_ENV_VARS = ("CUDA_INJECTION64_PATH",
"SMODEL_INJECTION_CONFIG")
_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 4ad200e912..ba6f30f3d7 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
test_environment_validation_is_not_process_cached(tmp_path, monkeypatch):
_iket_official.validate_official_environment()
[email protected](
+ ("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():
