This is an automated email from the ASF dual-hosted git repository.
tqchen 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 77a2ede9e9 [FIX][TIRx][CUDA] Support IKET 4.8 injection and trace
artifacts (#20554)
77a2ede9e9 is described below
commit 77a2ede9e914dbe1916e8e575b9e1dac56b34860
Author: Hongyi Jin <[email protected]>
AuthorDate: Mon Oct 5 21:45:10 2026 -0400
[FIX][TIRx][CUDA] Support IKET 4.8 injection and trace artifacts (#20554)
IKET 4.8.0 sets `IKET_INJECTION_CONFIG` in its tracker and capture
children. After #20477 removes the version checks, TVM still rejects
that environment in Python, and the C++ lowering path does not enable
annotations for plain `tvm.compile` calls. Recognize the new variable in
both paths, prefer it when both names are set, and retain the
`SMODEL_INJECTION_CONFIG` fallback.
Collect only `*.trace.json` as decoded traces so `run-iket.stats.json`
cannot be mistaken for a trace or hide missing output. Add regression
coverage for both environment names, precedence, nested injection and
statistics-only output, plus an `--arch` option for the GPU validation
workload.
Validation:
- Changed-file pre-commit checks pass on this PR branch, based on
`e292f6eec0`.
- GPU validation used this compatibility fix on `v0.27.0.post1`
(`7b5acb32`) plus #20477, with CuTe DSL / IKET 4.8.0 and CUDA/NVRTC 13.4
on Rubin (`sm_107a`): 89 IKET tests passed, and the real canonical trace
contract passed separately, including NativeDump and Extended payload
decoding.
- All 49 Rubin-supported TIRx-kernels configurations (32 FA4 and 17
FlashMLA) passed their numerical oracles without injection and in both
tracker/capture passes. Each produced decoded JSON and Perfetto output.
The matrix traced cluster 0 while checking the full numerical output;
separate FA4 and FlashMLA entry-point captures used the default trace
scope.
This PR contains only the additional compatibility fix, ported onto
current main. The full GPU matrix above was run on the explicitly stated
post1 + #20477 baseline, not rerun on this main-based branch. The legacy
external patch-log contract test was not run.
Companion TIRx-kernels entry-point and dependency fixes:
https://github.com/mlc-ai/TIRx-kernels/pull/5.
Co-authored-by: Codex <[email protected]>
---
python/tvm/backend/cuda/iket.py | 24 +++++++-----
src/backend/cuda/transforms/lower_iket.cc | 5 ++-
tests/python/tirx/iket/iket_profile_workload.py | 7 +++-
tests/python/tirx/iket/test_iket_orchestration.py | 33 +++++++++++++++-
tests/python/tirx/iket/test_iket_profiler.py | 48 ++++++++++++++++++-----
5 files changed, 95 insertions(+), 22 deletions(-)
diff --git a/python/tvm/backend/cuda/iket.py b/python/tvm/backend/cuda/iket.py
index 5f492b1ed4..b1fe5ed5b2 100644
--- a/python/tvm/backend/cuda/iket.py
+++ b/python/tvm/backend/cuda/iket.py
@@ -51,7 +51,8 @@ from tvm.script import tirx as T
_INJECTED_CHILD_ENABLE_ENV = "TVM_IKET_INJECTED_CHILD_ENABLE"
_POSTPROCESS_CHOICES = frozenset(("perfetto", "json", "html", "none", "all"))
-_INJECTION_ENV_VARS = ("CUDA_INJECTION64_PATH", "SMODEL_INJECTION_CONFIG")
+_INJECTION_CONFIG_ENV_VARS = ("IKET_INJECTION_CONFIG",
"SMODEL_INJECTION_CONFIG")
+_INJECTION_ENV_VARS = ("CUDA_INJECTION64_PATH", *_INJECTION_CONFIG_ENV_VARS)
_OUTPUT_TAIL_LINES = 100
_TERMINATION_GRACE_SECONDS = 5.0
@@ -196,19 +197,25 @@ def _validate_injection_environment() -> None:
if injection_path is None or not injection_path.is_file():
raise _profile_error("CUDA_INJECTION64_PATH was not supplied by
run-iket")
- config_value = os.environ.get("SMODEL_INJECTION_CONFIG")
+ # IKET 4.8 renamed the configuration variable. Prefer the current name
+ # when both are present, matching the injection library that run-iket uses.
+ config_name = next(
+ (name for name in _INJECTION_CONFIG_ENV_VARS if os.environ.get(name)),
+ _INJECTION_CONFIG_ENV_VARS[0],
+ )
+ config_value = os.environ.get(config_name)
config_path = Path(config_value) if config_value else None
if config_path is None or not config_path.is_file():
- raise _profile_error("SMODEL_INJECTION_CONFIG was not supplied by
run-iket")
+ raise _profile_error(f"{config_name} was not supplied by run-iket")
try:
config = json.loads(config_path.read_text(encoding="utf-8"))
except (OSError, UnicodeDecodeError, json.JSONDecodeError) as err:
- raise _profile_error("SMODEL_INJECTION_CONFIG is not valid JSON") from
err
+ raise _profile_error(f"{config_name} is not valid JSON") from err
if not isinstance(config, dict):
- raise _profile_error("SMODEL_INJECTION_CONFIG must contain a JSON
object")
+ raise _profile_error(f"{config_name} must contain a JSON object")
tool_name = config.get("toolName")
if tool_name not in ("tracker", "iket"):
- raise _profile_error("SMODEL_INJECTION_CONFIG was not generated by
run-iket profile")
+ raise _profile_error(f"{config_name} was not generated by run-iket
profile")
if tool_name == "tracker":
return
tool_config = config.get("toolConfig")
@@ -557,9 +564,8 @@ def _run_process(
def _collect_artifacts(
output_dir: Path,
) -> tuple[tuple[Path, ...], tuple[Path, ...], tuple[Path, ...]]:
- root_json = {path for path in output_dir.glob("*.json") if path.is_file()}
- nested_trace_json = {path for path in output_dir.rglob("*.trace.json") if
path.is_file()}
- json_traces = tuple(sorted(root_json | nested_trace_json))
+ # run-iket can also write statistics/configuration JSON. Those are not
traces.
+ json_traces = tuple(sorted(path for path in
output_dir.rglob("*.trace.json") if path.is_file()))
perfetto_traces = tuple(
sorted(path for path in output_dir.rglob("*.pftrace") if
path.is_file())
)
diff --git a/src/backend/cuda/transforms/lower_iket.cc
b/src/backend/cuda/transforms/lower_iket.cc
index 49fe69b35a..64138c9f8e 100644
--- a/src/backend/cuda/transforms/lower_iket.cc
+++ b/src/backend/cuda/transforms/lower_iket.cc
@@ -1210,7 +1210,10 @@ bool IketEnabled(const IRModule& module) {
if (module->HasNonzeroAttr("tirx.iket.enabled")) return true;
const char* child_enable = std::getenv("TVM_IKET_INJECTED_CHILD_ENABLE");
const char* injection = std::getenv("CUDA_INJECTION64_PATH");
- const char* injection_config = std::getenv("SMODEL_INJECTION_CONFIG");
+ const char* injection_config = std::getenv("IKET_INJECTION_CONFIG");
+ if (!injection_config || injection_config[0] == '\0') {
+ injection_config = std::getenv("SMODEL_INJECTION_CONFIG");
+ }
return child_enable && std::string(child_enable) == "1" && injection &&
injection[0] != '\0' &&
injection_config && injection_config[0] != '\0';
}
diff --git a/tests/python/tirx/iket/iket_profile_workload.py
b/tests/python/tirx/iket/iket_profile_workload.py
index ab178f3528..37638fefaf 100644
--- a/tests/python/tirx/iket/iket_profile_workload.py
+++ b/tests/python/tirx/iket/iket_profile_workload.py
@@ -122,6 +122,7 @@ def _parse_args():
parser.add_argument("--keep", action=argparse.BooleanOptionalAction,
default=False)
parser.add_argument("--timeout", type=float, default=600.0)
parser.add_argument("--max-ts-cnt-per-warp", type=int, default=None)
+ parser.add_argument("--arch", default="sm_100a")
parser.add_argument(
"--fail-capture",
action="store_true",
@@ -131,7 +132,9 @@ def _parse_args():
def _injection_tool_name():
- config_path = os.environ.get("SMODEL_INJECTION_CONFIG")
+ config_path = os.environ.get("IKET_INJECTION_CONFIG") or os.environ.get(
+ "SMODEL_INJECTION_CONFIG"
+ )
if not config_path:
return None
return
json.loads(Path(config_path).read_text(encoding="utf-8")).get("toolName")
@@ -140,7 +143,7 @@ def _injection_tool_name():
def _profile_workload(args):
if args.fail_capture and _injection_tool_name() == "iket":
raise RuntimeError("intentional capture-only IKET workload failure")
- target = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"})
+ target = tvm.target.Target({"kind": "cuda", "arch": args.arch})
workloads = (
(canonical_iket_workload, 1),
(native_payload_workload, 2),
diff --git a/tests/python/tirx/iket/test_iket_orchestration.py
b/tests/python/tirx/iket/test_iket_orchestration.py
index 86d74af50b..dd3c6fe19e 100644
--- a/tests/python/tirx/iket/test_iket_orchestration.py
+++ b/tests/python/tirx/iket/test_iket_orchestration.py
@@ -67,7 +67,10 @@ def test_public_namespace_and_old_bench_path_is_absent():
assert not hasattr(bench, "IketProfiler")
[email protected]("injection_variable", ["CUDA_INJECTION64_PATH",
"SMODEL_INJECTION_CONFIG"])
[email protected](
+ "injection_variable",
+ ["CUDA_INJECTION64_PATH", "IKET_INJECTION_CONFIG",
"SMODEL_INJECTION_CONFIG"],
+)
def test_profile_rejects_nested_injection_before_validation(
injection_variable, tmp_path, monkeypatch
):
@@ -200,6 +203,33 @@ def
test_multi_process_singular_properties_point_to_plural(tmp_path, monkeypatch
_ = result.html
+def test_statistics_json_is_not_a_trace(tmp_path, monkeypatch):
+ def run_process(argv, **_kwargs):
+ staging = Path(argv[argv.index("--output-dir") + 1])
+ _write_artifacts(staging, "json")
+ (staging / "run-iket.stats.json").write_text('{"phases": {}}',
encoding="utf-8")
+
+ monkeypatch.setattr(iket, "_validate_official_installation", lambda:
"run-iket")
+ monkeypatch.setattr(iket, "_run_process", run_process)
+ result = iket.profile(
+ ("python", "workload.py"), output_dir=tmp_path / "result",
postprocess="json"
+ )
+ assert result.trace == {"launches": [{"index": 0}]}
+ assert len(result.json_traces) == 1
+ assert (result.output_dir / "run-iket.stats.json").is_file()
+
+
+def test_statistics_json_does_not_satisfy_requested_trace(tmp_path,
monkeypatch):
+ def run_process(argv, **_kwargs):
+ staging = Path(argv[argv.index("--output-dir") + 1])
+ (staging / "run-iket.stats.json").write_text('{"phases": {}}',
encoding="utf-8")
+
+ monkeypatch.setattr(iket, "_validate_official_installation", lambda:
"run-iket")
+ monkeypatch.setattr(iket, "_run_process", run_process)
+ with pytest.raises(iket.IketProfileError, match="JSON trace"):
+ iket.profile(("python", "workload.py"), output_dir=tmp_path /
"result", postprocess="json")
+
+
def test_missing_artifact_rolls_back_existing_output(tmp_path, monkeypatch):
target = tmp_path / "result"
target.mkdir()
@@ -435,6 +465,7 @@ def
test_imported_callable_is_rejected_before_profile_starts(tmp_path, monkeypat
monkeypatch.delenv("CUDA_INJECTION64_PATH", raising=False)
monkeypatch.delenv("SMODEL_INJECTION_CONFIG", raising=False)
+ monkeypatch.delenv("IKET_INJECTION_CONFIG", raising=False)
monkeypatch.setattr(iket, "profile", should_not_start)
with pytest.raises(ValueError, match=r"iket\.profile\(command\)"):
iket.run(test_imported_callable_is_rejected_before_profile_starts,
output_dir=tmp_path)
diff --git a/tests/python/tirx/iket/test_iket_profiler.py
b/tests/python/tirx/iket/test_iket_profiler.py
index 87067d9824..71f2692dd5 100644
--- a/tests/python/tirx/iket/test_iket_profiler.py
+++ b/tests/python/tirx/iket/test_iket_profiler.py
@@ -433,11 +433,14 @@ def test_regular_lowering_strips_annotations_and_tokens():
assert "iket" not in sources[1].lower()
-def test_verified_injected_child_automatically_enables_plain_jit(monkeypatch):
[email protected]("config_env", ["IKET_INJECTION_CONFIG",
"SMODEL_INJECTION_CONFIG"])
+def test_verified_injected_child_automatically_enables_plain_jit(monkeypatch,
config_env):
+ monkeypatch.delenv("IKET_INJECTION_CONFIG", raising=False)
+ monkeypatch.delenv("SMODEL_INJECTION_CONFIG", raising=False)
monkeypatch.setenv("TVM_IKET_INJECTED_CHILD_ENABLE", "1")
monkeypatch.delenv("TVM_IKET_OFFICIAL_PROFILE", raising=False)
monkeypatch.setenv("CUDA_INJECTION64_PATH",
"/verified/libsmodel_injection.so")
- monkeypatch.setenv("SMODEL_INJECTION_CONFIG", "/verified/config.json")
+ monkeypatch.setenv(config_env, "/verified/config.json")
executable = tvm.compile(tvm.IRModule({"main": serial_a}), target=TARGET,
tir_pipeline="tirx")
source = _cuda_source(executable)
assert "__iket_evt_decl_a_1_attrs" in source
@@ -445,17 +448,20 @@ def
test_verified_injected_child_automatically_enables_plain_jit(monkeypatch):
@pytest.mark.parametrize(
"missing_env",
- ("CUDA_INJECTION64_PATH", "SMODEL_INJECTION_CONFIG",
"TVM_IKET_INJECTED_CHILD_ENABLE"),
+ ("CUDA_INJECTION64_PATH", "config", "TVM_IKET_INJECTED_CHILD_ENABLE"),
)
-def test_injected_child_auto_enable_remains_fail_closed(monkeypatch,
missing_env):
[email protected]("config_env", ["IKET_INJECTION_CONFIG",
"SMODEL_INJECTION_CONFIG"])
+def test_injected_child_auto_enable_remains_fail_closed(monkeypatch,
missing_env, config_env):
+ monkeypatch.delenv("IKET_INJECTION_CONFIG", raising=False)
+ monkeypatch.delenv("SMODEL_INJECTION_CONFIG", raising=False)
values = {
"TVM_IKET_INJECTED_CHILD_ENABLE": "1",
"CUDA_INJECTION64_PATH": "/verified/libsmodel_injection.so",
- "SMODEL_INJECTION_CONFIG": "/verified/config.json",
+ config_env: "/verified/config.json",
}
for name, value in values.items():
monkeypatch.setenv(name, value)
- monkeypatch.delenv(missing_env)
+ monkeypatch.delenv(config_env if missing_env == "config" else missing_env)
executable = tvm.compile(tvm.IRModule({"main": serial_a}), target=TARGET,
tir_pipeline="tirx")
assert "iket" not in _cuda_source(executable).lower()
@@ -864,14 +870,17 @@ def test_nvrtc_load_error(monkeypatch, error_type):
_iket_official._validate_nvrtc_available() # pylint:
disable=protected-access
-def test_injection_environment_accepts_run_iket_two_passes(tmp_path,
monkeypatch):
[email protected]("config_env", ["IKET_INJECTION_CONFIG",
"SMODEL_INJECTION_CONFIG"])
+def test_injection_environment_accepts_run_iket_two_passes(tmp_path,
monkeypatch, config_env):
from tvm.tirx.cuda import iket as _iket_official
- injection = tmp_path / "libsmodel_injection.so"
+ monkeypatch.delenv("IKET_INJECTION_CONFIG", raising=False)
+ monkeypatch.delenv("SMODEL_INJECTION_CONFIG", raising=False)
+ injection = tmp_path / "libiket_injection.so"
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))
+ monkeypatch.setenv(config_env, str(config_path))
config_path.write_text(json.dumps({"toolName": "tracker", "toolConfig":
{}}))
_iket_official._validate_injection_environment() # pylint:
disable=protected-access
@@ -893,6 +902,27 @@ def
test_injection_environment_accepts_run_iket_two_passes(tmp_path, monkeypatch
_iket_official._validate_injection_environment() # pylint:
disable=protected-access
+def test_injection_environment_prefers_current_config(tmp_path, monkeypatch):
+ from tvm.tirx.cuda import iket as _iket_official
+
+ injection = tmp_path / "libiket_injection.so"
+ injection.touch()
+ current = tmp_path / "current.json"
+ legacy = tmp_path / "legacy.json"
+ current.write_text(json.dumps({"toolName": "tracker"}), encoding="utf-8")
+ legacy.write_text("invalid", encoding="utf-8")
+ monkeypatch.setenv("CUDA_INJECTION64_PATH", str(injection))
+ monkeypatch.setenv("IKET_INJECTION_CONFIG", str(current))
+ monkeypatch.setenv("SMODEL_INJECTION_CONFIG", str(legacy))
+ _iket_official._validate_injection_environment()
+ current.write_text("invalid", encoding="utf-8")
+ legacy.write_text(json.dumps({"toolName": "tracker"}), encoding="utf-8")
+ with pytest.raises(RuntimeError, match="IKET_INJECTION_CONFIG is not valid
JSON"):
+ _iket_official._validate_injection_environment()
+ monkeypatch.setenv("IKET_INJECTION_CONFIG", "")
+ _iket_official._validate_injection_environment()
+
+
def test_cutlass_4_6_0_oracle_manifest_integrity():
oracle = json.loads(ORACLE_PATH.read_text(encoding="utf-8"))
assert oracle["schema_version"] == 3
