| |
| from __future__ import annotations |
|
|
| import argparse |
| import ctypes |
| import hashlib |
| import json |
| import os |
| import subprocess |
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
| import tensorrt as trt |
|
|
|
|
| HERE = Path(__file__).resolve().parent |
| ARTIFACTS = HERE / "artifacts" |
| CONTROL = ARTIFACTS / "control_msda_left.engine" |
| MALICIOUS = ARTIFACTS / "malicious_msda_right.engine" |
|
|
|
|
| def sha256(path: Path) -> str: |
| h = hashlib.sha256() |
| with path.open("rb") as f: |
| for chunk in iter(lambda: f.read(1024 * 1024), b""): |
| h.update(chunk) |
| return h.hexdigest() |
|
|
|
|
| def load_cudart(): |
| cudart = ctypes.CDLL("libcudart.so") |
| cudart.cudaSetDevice.argtypes = [ctypes.c_int] |
| cudart.cudaSetDevice.restype = ctypes.c_int |
| cudart.cudaMalloc.argtypes = [ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t] |
| cudart.cudaMalloc.restype = ctypes.c_int |
| cudart.cudaFree.argtypes = [ctypes.c_void_p] |
| cudart.cudaFree.restype = ctypes.c_int |
| cudart.cudaMemcpy.argtypes = [ |
| ctypes.c_void_p, |
| ctypes.c_void_p, |
| ctypes.c_size_t, |
| ctypes.c_int, |
| ] |
| cudart.cudaMemcpy.restype = ctypes.c_int |
| cudart.cudaMemset.argtypes = [ctypes.c_void_p, ctypes.c_int, ctypes.c_size_t] |
| cudart.cudaMemset.restype = ctypes.c_int |
| cudart.cudaDeviceSynchronize.restype = ctypes.c_int |
| return cudart |
|
|
|
|
| def tensor_count(shape: list[int]) -> int: |
| count = 1 |
| for dim in shape: |
| if dim < 0: |
| raise RuntimeError(f"dynamic tensor shape is not bound: {shape}") |
| count *= dim |
| return count |
|
|
|
|
| def get_creator() -> trt.IPluginCreator: |
| logger = trt.Logger(trt.Logger.ERROR) |
| trt.init_libnvinfer_plugins(logger, "") |
| for creator in trt.get_plugin_registry().all_creators: |
| if ( |
| creator.name == "MultiscaleDeformableAttnPlugin_TRT" |
| and creator.plugin_version == "1" |
| and type(creator).__name__ == "IPluginCreator" |
| ): |
| return creator |
| raise RuntimeError("MultiscaleDeformableAttnPlugin_TRT v1 creator not found") |
|
|
|
|
| def build_engine(path: Path, sampling_xy: tuple[float, float]) -> None: |
| logger = trt.Logger(trt.Logger.ERROR) |
| creator = get_creator() |
| plugin = creator.create_plugin("msda_gate", trt.PluginFieldCollection([])) |
| if plugin is None: |
| raise RuntimeError("create_plugin returned None") |
|
|
| builder = trt.Builder(logger) |
| network = builder.create_network(0) |
| config = builder.create_builder_config() |
|
|
| value = network.add_input("value", trt.float32, (1, 2, 1, 1)) |
| spatial_shapes = network.add_constant( |
| (1, 2), np.array([1, 2], dtype=np.int32) |
| ).get_output(0) |
| level_start_index = network.add_constant( |
| (1,), np.array([0], dtype=np.int32) |
| ).get_output(0) |
| sampling_locations = network.add_constant( |
| (1, 1, 1, 1, 1, 2), np.array(sampling_xy, dtype=np.float32) |
| ).get_output(0) |
| attention_weights = network.add_constant( |
| (1, 1, 1, 1, 1), np.array([1.0], dtype=np.float32) |
| ).get_output(0) |
|
|
| layer = network.add_plugin_v2( |
| [value, spatial_shapes, level_start_index, sampling_locations, attention_weights], |
| plugin, |
| ) |
| if layer is None: |
| raise RuntimeError("add_plugin_v2 returned None") |
| layer.get_output(0).name = "attn_out" |
| network.mark_output(layer.get_output(0)) |
|
|
| plan = builder.build_serialized_network(network, config) |
| if plan is None: |
| raise RuntimeError("TensorRT failed to build MultiscaleDeformableAttn engine") |
| path.parent.mkdir(parents=True, exist_ok=True) |
| path.write_bytes(bytes(plan)) |
|
|
|
|
| def build_artifacts() -> None: |
| ARTIFACTS.mkdir(parents=True, exist_ok=True) |
| build_engine(CONTROL, (0.25, 0.5)) |
| build_engine(MALICIOUS, (0.75, 0.5)) |
|
|
|
|
| def run_engine(path: Path, values: list[float], gpu: int | None = None) -> dict: |
| cuda = load_cudart() |
| if gpu is not None: |
| rc = cuda.cudaSetDevice(gpu) |
| if rc != 0: |
| raise RuntimeError(f"cudaSetDevice({gpu}) failed rc={rc}") |
|
|
| logger = trt.Logger(trt.Logger.ERROR) |
| trt.init_libnvinfer_plugins(logger, "") |
| runtime = trt.Runtime(logger) |
| engine = runtime.deserialize_cuda_engine(path.read_bytes()) |
| if engine is None: |
| raise RuntimeError(f"deserialize_cuda_engine returned None for {path}") |
| ctx = engine.create_execution_context() |
| if ctx is None: |
| raise RuntimeError("create_execution_context returned None") |
|
|
| ptrs: dict[str, tuple[ctypes.c_void_p, int, list[int], str]] = {} |
| try: |
| tensor_meta = {} |
| for i in range(engine.num_io_tensors): |
| name = engine.get_tensor_name(i) |
| shape = list(engine.get_tensor_shape(name)) |
| count = tensor_count(shape) |
| nbytes = max(count * 4, 4) |
| ptr = ctypes.c_void_p() |
| if cuda.cudaMalloc(ctypes.byref(ptr), nbytes) != 0: |
| raise RuntimeError(f"cudaMalloc failed for {name}") |
| if cuda.cudaMemset(ptr, 0, nbytes) != 0: |
| raise RuntimeError(f"cudaMemset failed for {name}") |
| if not ctx.set_tensor_address(name, ptr.value): |
| raise RuntimeError(f"set_tensor_address failed for {name}") |
| mode = str(engine.get_tensor_mode(name)) |
| ptrs[name] = (ptr, nbytes, shape, mode) |
| tensor_meta[name] = { |
| "shape": shape, |
| "dtype": str(engine.get_tensor_dtype(name)), |
| "mode": mode, |
| } |
|
|
| arr = (ctypes.c_float * len(values))(*values) |
| upload_rc = cuda.cudaMemcpy( |
| ptrs["value"][0], ctypes.cast(arr, ctypes.c_void_p), len(values) * 4, 1 |
| ) |
| if upload_rc != 0: |
| raise RuntimeError(f"cudaMemcpy input upload failed rc={upload_rc}") |
|
|
| infer_shapes = ctx.infer_shapes() |
| execute_ok = bool(ctx.execute_async_v3(0)) |
| sync_rc = cuda.cudaDeviceSynchronize() |
| outputs = {} |
| for name, (ptr, nbytes, shape, mode) in ptrs.items(): |
| if "OUTPUT" not in mode: |
| continue |
| count = nbytes // 4 |
| host = (ctypes.c_float * count)() |
| copy_rc = cuda.cudaMemcpy(ctypes.byref(host), ptr, nbytes, 2) |
| outputs[name] = { |
| "shape": shape, |
| "copy_rc": copy_rc, |
| "values": [float(x) for x in host], |
| } |
|
|
| return { |
| "path": str(path), |
| "sha256": sha256(path), |
| "input_values": [float(x) for x in values], |
| "tensor_meta": tensor_meta, |
| "infer_shapes": infer_shapes, |
| "execute_ok": execute_ok, |
| "sync_rc": sync_rc, |
| "outputs": outputs, |
| "output_signature": [ |
| value |
| for output in outputs.values() |
| for value in output.get("values", []) |
| ], |
| } |
| finally: |
| for ptr, _, _, _ in ptrs.values(): |
| if ptr.value: |
| cuda.cudaFree(ptr) |
|
|
|
|
| def run(cmd: list[str], timeout: int = 120) -> dict: |
| proc = subprocess.run(cmd, text=True, capture_output=True, timeout=timeout) |
| return { |
| "returncode": proc.returncode, |
| "stdout": proc.stdout.strip(), |
| "stderr_tail": proc.stderr[-3000:], |
| } |
|
|
|
|
| def modelscan(path: Path) -> dict: |
| modelscan_env = os.environ.get("MODELSCAN_BIN") |
| modelscan_bin = Path(modelscan_env) if modelscan_env else None |
| if not modelscan_bin or not modelscan_bin.exists(): |
| modelscan_bin = Path(sys.executable).with_name("modelscan") |
| if not modelscan_bin.exists(): |
| modelscan_bin = Path("modelscan") |
| result = run([str(modelscan_bin), "-p", str(path), "--show-skipped"]) |
| combined = result["stdout"] + "\n" + result["stderr_tail"] |
| result["binary"] = str(modelscan_bin) |
| result["clean"] = "No issues found" in combined |
| result["skipped"] = "Model Scan did not scan file" in combined |
| return result |
|
|
|
|
| def summarize(control: Path, malicious: Path, gpu: int | None = None) -> dict: |
| benign = [7.0, 7.0] |
| trigger = [10.0, 99.0] |
| summary = { |
| "format": "TensorRT (.engine / .trt / .mytrtfile) - NVIDIA", |
| "tensorrt": trt.__version__, |
| "plugin": "MultiscaleDeformableAttnPlugin_TRT v1", |
| "hidden_constants": { |
| "control_sampling_locations_xy": [0.25, 0.5], |
| "malicious_sampling_locations_xy": [0.75, 0.5], |
| "spatial_shapes": [1, 2], |
| "level_start_index": [0], |
| "attention_weights": [1.0], |
| }, |
| "control": str(control), |
| "malicious": str(malicious), |
| "sha256": { |
| "control": sha256(control), |
| "malicious": sha256(malicious), |
| }, |
| "runs": { |
| "control_benign": run_engine(control, benign, gpu), |
| "malicious_benign": run_engine(malicious, benign, gpu), |
| "control_trigger": run_engine(control, trigger, gpu), |
| "malicious_trigger": run_engine(malicious, trigger, gpu), |
| }, |
| "modelscan_malicious": modelscan(malicious), |
| } |
| summary["impact"] = { |
| "same_output_shape": ( |
| summary["runs"]["control_trigger"]["tensor_meta"]["attn_out"]["shape"] |
| == summary["runs"]["malicious_trigger"]["tensor_meta"]["attn_out"]["shape"] |
| ), |
| "both_execute_ok": ( |
| summary["runs"]["control_trigger"]["execute_ok"] |
| and summary["runs"]["malicious_trigger"]["execute_ok"] |
| ), |
| "sync_ok": ( |
| summary["runs"]["control_trigger"]["sync_rc"] == 0 |
| and summary["runs"]["malicious_trigger"]["sync_rc"] == 0 |
| ), |
| "benign_outputs_match": ( |
| summary["runs"]["control_benign"]["output_signature"] |
| == summary["runs"]["malicious_benign"]["output_signature"] |
| ), |
| "trigger_output_changed": ( |
| summary["runs"]["control_trigger"]["output_signature"] |
| != summary["runs"]["malicious_trigger"]["output_signature"] |
| ), |
| "control_trigger_output": summary["runs"]["control_trigger"]["output_signature"], |
| "malicious_trigger_output": summary["runs"]["malicious_trigger"]["output_signature"], |
| } |
| return summary |
|
|
|
|
| def main() -> int: |
| ap = argparse.ArgumentParser( |
| description="Build and reproduce a MultiscaleDeformableAttn TensorRT trigger backdoor." |
| ) |
| ap.add_argument("--build", action="store_true") |
| ap.add_argument("--control", type=Path, default=CONTROL) |
| ap.add_argument("--malicious", type=Path, default=MALICIOUS) |
| ap.add_argument("--gpu", type=int, default=None) |
| ap.add_argument("--out", type=Path, default=ARTIFACTS / "repro_summary.json") |
| args = ap.parse_args() |
|
|
| if args.build: |
| build_artifacts() |
|
|
| summary = summarize(args.control.resolve(), args.malicious.resolve(), args.gpu) |
| args.out.parent.mkdir(parents=True, exist_ok=True) |
| args.out.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8") |
| print(json.dumps(summary, indent=2)) |
|
|
| impact = summary["impact"] |
| if not ( |
| impact["same_output_shape"] |
| and impact["both_execute_ok"] |
| and impact["sync_ok"] |
| and impact["benign_outputs_match"] |
| and impact["trigger_output_changed"] |
| and summary["modelscan_malicious"]["clean"] |
| ): |
| raise SystemExit("FAIL: expected MultiscaleDeformableAttn trigger did not reproduce") |
| print("PASS: malicious MultiscaleDeformableAttn engine reroutes trigger output while modelscan reports clean") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
