diff --git a/benchmarks/performance/baselines/task_reference.py b/benchmarks/performance/baselines/task_reference.py index d88019cd0e..fdba984560 100644 --- a/benchmarks/performance/baselines/task_reference.py +++ b/benchmarks/performance/baselines/task_reference.py @@ -573,7 +573,7 @@ def _load_asr( max_new_tokens = int(request.get("max_new_tokens", 100)) device = torch.device("cuda") - if arguments.family in {"canary", "nemotron_speech_streaming"}: + if arguments.family in {"canary", "nemotron_speech_streaming", "timm_repvgg"}: from tools.validation.engine import _transcription_text model = _load_nemo_asr_reference_model(arguments, device=device).eval().to(device) diff --git a/benchmarks/performance/baselines/timing_contracts.py b/benchmarks/performance/baselines/timing_contracts.py index 3ce193ff44..0ddebe4d79 100644 --- a/benchmarks/performance/baselines/timing_contracts.py +++ b/benchmarks/performance/baselines/timing_contracts.py @@ -25,6 +25,7 @@ "sam3", "segformer", "timesfm", + "timm_repvgg", "timm_resnet", "timm_vgg", "timm_vit", diff --git a/benchmarks/performance/release.yaml b/benchmarks/performance/release.yaml index bb5f41292f..3fb787a467 100644 --- a/benchmarks/performance/release.yaml +++ b/benchmarks/performance/release.yaml @@ -974,6 +974,19 @@ entries: reference_backend: hf_transformers timing_scope: task-model-call-wall input_preparation_included: false + - id: timm_repvgg.classify + family: timm_repvgg + operation: classify + model: repvgg-a2-rvgg-in1k + workload: + testcase: repvgg-a2-rvgg-in1k + baseline: + runner: task-reference + adapter: hf-transformers-vision + mode: hf-eager + reference_backend: hf_transformers + timing_scope: task-model-call-wall + input_preparation_included: false - id: timm_vgg.classify family: timm_vgg operation: classify diff --git a/python/tensorrt_model_connect/families/timm_repvgg/MODEL.toml b/python/tensorrt_model_connect/families/timm_repvgg/MODEL.toml new file mode 100644 index 0000000000..b4ee0d7da2 --- /dev/null +++ b/python/tensorrt_model_connect/families/timm_repvgg/MODEL.toml @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +id = "timm_repvgg" +plugin = "timm_repvgg" +module = "plugin" +python_profile_specs = [ + "timm_repvgg_reference|families/timm_repvgg/python_profile_requirements/timm_repvgg_reference.lock.txt|families/timm_repvgg/python_profile_verify.py|true", +] +default_execution_profiles = [ + "reference|timm_repvgg_reference", +] +aliases = [ + "timm_repvgg", + "repvgg", + "repvgg_a2", +] +prefixes = [ + "timm_repvgg", + "repvgg", +] diff --git a/python/tensorrt_model_connect/families/timm_repvgg/__init__.py b/python/tensorrt_model_connect/families/timm_repvgg/__init__.py new file mode 100644 index 0000000000..f1c2efb4b0 --- /dev/null +++ b/python/tensorrt_model_connect/families/timm_repvgg/__init__.py @@ -0,0 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from .plugin import plugin + +__all__ = ["plugin"] diff --git a/python/tensorrt_model_connect/families/timm_repvgg/config.py b/python/tensorrt_model_connect/families/timm_repvgg/config.py new file mode 100644 index 0000000000..d062ab0237 --- /dev/null +++ b/python/tensorrt_model_connect/families/timm_repvgg/config.py @@ -0,0 +1,239 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""ModelConfig — parse HF config.json into a typed dataclass.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from pathlib import Path + + +@dataclass +class ModelConfig: + """Parsed model architecture from HF config.json.""" + + model_type: str = "" + architectures: list[str] = field(default_factory=list) + vocab_size: int = 0 + hidden_size: int = 0 + intermediate_size: int = 0 + num_hidden_layers: int = 0 + num_attention_heads: int = 1 + num_key_value_heads: int = 1 + rms_norm_eps: float = 1e-5 + rope_theta: float = 10000.0 + bos_token_id: int = -1 + eos_token_id: int = -1 + pad_token_id: int = -1 + tie_word_embeddings: bool = False + max_position_embeddings: int = 8192 + hidden_act: str = "" + + # Explicit head_dim from config.json (0 = not set, fall back to computed). + _head_dim: int = 0 + + # Raw JSON dict for family-specific fields + raw: dict = field(default_factory=dict, repr=False) + + @property + def head_dim(self) -> int: + if self._head_dim > 0: + return self._head_dim + if self.num_attention_heads <= 0: + return 0 + return self.hidden_size // self.num_attention_heads + + @property + def attention_size(self) -> int: + return self.num_attention_heads * self.head_dim + + @staticmethod + def from_json(text: str) -> ModelConfig: + d = json.loads(text) + + # Some multimodal configs nest decoder fields under "text_config". + # Merge text_config into top level so standard key lookup works. + # Preserve top-level model_type and architectures (these identify the + # top-level model, not the nested decoder). + original_raw = d + text_config = d.get("text_config") + if text_config and isinstance(text_config, dict): + top_model_type = d.get("model_type") + top_architectures = d.get("architectures") + merged = {**d, **text_config} + if top_model_type: + merged["model_type"] = top_model_type + if top_architectures: + merged["architectures"] = top_architectures + d = merged + + # Some multimodal configs nest the language decoder config under + # "language_config". Merge into top level like text_config. + if not d.get("hidden_size"): + lang_config = d.get("language_config") + if isinstance(lang_config, dict): + top_model_type = d.get("model_type") + top_architectures = d.get("architectures") + top_vision_config = d.get("vision_config") + merged = {**d, **lang_config} + if top_model_type: + merged["model_type"] = top_model_type + if top_architectures: + merged["architectures"] = top_architectures + if top_vision_config: + merged["vision_config"] = top_vision_config + d = merged + + # Some multimodal configs nest LLM config under "llm_config". + # Merge into top level like text_config, preserving top-level + # model_type, architectures, and vision_config. + if not d.get("hidden_size"): + llm_config = d.get("llm_config") + if isinstance(llm_config, dict): + top_model_type = d.get("model_type") + top_architectures = d.get("architectures") + top_vision_config = d.get("vision_config") + merged = {**d, **llm_config} + if top_model_type: + merged["model_type"] = top_model_type + if top_architectures: + merged["architectures"] = top_architectures + if top_vision_config: + merged["vision_config"] = top_vision_config + d = merged + + # Some multimodal audio/text configs nest the primary decoder config + # under thinker_config.text_config. If top-level hidden_size is + # still missing after the text_config merge above, look there. + if not d.get("hidden_size"): + thinker_cfg = d.get("thinker_config") + if isinstance(thinker_cfg, dict): + thinker_text = thinker_cfg.get("text_config") + if isinstance(thinker_text, dict): + top_model_type = d.get("model_type") + top_architectures = d.get("architectures") + merged = {**d, **thinker_text} + if top_model_type: + merged["model_type"] = top_model_type + if top_architectures: + merged["architectures"] = top_architectures + # Also propagate vision_config from thinker_config + # so VL pipelines can find it. + if "vision_config" not in merged and "vision_config" in thinker_cfg: + merged["vision_config"] = thinker_cfg["vision_config"] + d = merged + + # Handle non-standard config key names: + # GPT-2: n_embd, n_head, n_layer, n_inner + # XGLM/Bloom: d_model, attention_heads, num_layers, ffn_dim + # DistilBERT: dim, n_heads, n_layers, hidden_dim + hidden_size = ( + d.get("hidden_size", 0) + or d.get("n_embd", 0) + or d.get("d_model", 0) + or d.get("n_embed", 0) + or d.get("dim", 0) + ) + num_heads = ( + d.get("num_attention_heads", 0) + or d.get("n_head", 0) + or d.get("attention_heads", 0) + or d.get("num_heads", 0) + or d.get("n_heads", 0) + or d.get("decoder_attention_heads", 0) + or 1 + ) + num_layers = ( + d.get("num_hidden_layers", 0) + or d.get("n_layer", 0) + or d.get("num_layers", 0) + or d.get("n_layers", 0) + ) + intermediate = ( + d.get("intermediate_size", 0) + or d.get("n_inner", 0) + or d.get("ffn_dim", 0) + or d.get("hidden_dim", 0) + or hidden_size * 4 + ) + + # Norm epsilon: try rms_norm_eps, then layer_norm_epsilon, then + # layer_norm_eps, then norm_epsilon, then norm_eps. + eps = ( + d.get("rms_norm_eps") + or d.get("layer_norm_epsilon") + or d.get("layer_norm_eps") + or d.get("norm_epsilon") + or d.get("norm_eps") + or 1e-5 + ) + + # rope_theta: check top-level first, then rope_parameters dict + # (some model configs store it there), + # then rope_scaling dict. + rope_theta = d.get("rope_theta", None) + if rope_theta is None: + rope_params = d.get("rope_parameters") + if isinstance(rope_params, dict): + rope_theta = rope_params.get("rope_theta", 10000.0) + else: + rope_scaling = d.get("rope_scaling") + if isinstance(rope_scaling, dict): + rope_theta = rope_scaling.get("rope_theta", 10000.0) + else: + rope_theta = 10000.0 + rope_theta = float(rope_theta) + + architecture = d.get("architecture", "") + architectures = d.get("architectures", []) + if not architectures and architecture: + architectures = [architecture] + + return ModelConfig( + model_type=d.get("model_type", "") or architecture, + architectures=architectures, + vocab_size=d.get("vocab_size", 0), + hidden_size=hidden_size or d.get("num_features", 0), + intermediate_size=intermediate, + num_hidden_layers=num_layers, + num_attention_heads=num_heads, + num_key_value_heads=d.get("num_key_value_heads", num_heads), + rms_norm_eps=eps, + rope_theta=rope_theta, + bos_token_id=d.get("bos_token_id", -1) or -1, + eos_token_id=d.get("eos_token_id", -1) or -1, + pad_token_id=d.get("pad_token_id", -1) or -1, + tie_word_embeddings=d.get("tie_word_embeddings", False), + max_position_embeddings=d.get("max_position_embeddings", d.get("n_positions", 8192)), + hidden_act=d.get("hidden_act", "") or d.get("activation_function", ""), + _head_dim=d.get("head_dim", 0), + raw=original_raw, + ) + + @classmethod + def create_tiny(cls, model_type: str, **overrides) -> "ModelConfig": + """Create a minimal ModelConfig for testing (2 layers, hidden=16, vocab=32).""" + defaults = { + "model_type": model_type, + "vocab_size": 32, + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "rms_norm_eps": 1e-6, + "rope_theta": 10000.0, + "max_position_embeddings": 128, + } + defaults.update(overrides) + return cls.from_json(json.dumps(defaults)) + + @staticmethod + def from_dir(model_dir: str | Path) -> ModelConfig: + model_path = Path(model_dir) + config_path = model_path / "config.json" + if config_path.exists(): + return ModelConfig.from_json(config_path.read_text()) + return ModelConfig.from_json(config_path.read_text()) diff --git a/python/tensorrt_model_connect/families/timm_repvgg/model/__init__.py b/python/tensorrt_model_connect/families/timm_repvgg/model/__init__.py new file mode 100644 index 0000000000..a1947c32d7 --- /dev/null +++ b/python/tensorrt_model_connect/families/timm_repvgg/model/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Family-owned TensorRT model construction components.""" diff --git a/python/tensorrt_model_connect/families/timm_repvgg/model/model.py b/python/tensorrt_model_connect/families/timm_repvgg/model/model.py new file mode 100644 index 0000000000..719dd81054 --- /dev/null +++ b/python/tensorrt_model_connect/families/timm_repvgg/model/model.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""TensorRT graph builders for timm RepVGG classifiers. + +RepVGG trains as a multi-branch block and is meant to run as a plain 3x3 +convolution stack. The branches are fused on the host in the weight loader, so +the graph here needs nothing beyond convolution, ReLU, pooling, and a linear +head. +""" + +from __future__ import annotations + +import numpy as np +from tensorrt_model_connect import trt_compat + +trt = trt_compat.get_trt() + + +def add_conv2d( + network, + inp, + weight: np.ndarray, + bias: np.ndarray | None, + out_channels: int, + kernel_size: tuple[int, int], + stride: tuple[int, int] = (1, 1), + padding: tuple[int, int] = (0, 0), + groups: int = 1, + dtype: np.dtype = np.float32, +): + """2D convolution wrapper. + + Input: [N, C_in, H, W] + Weight: [C_out, C_in/groups, kH, kW] + Output: [N, C_out, H', W'] + + timm ResNets always fold the convolution bias into the following batch + norm, so `bias` is None throughout this family; the parameter is kept so + the signature matches the other families that own this helper. + """ + conv_w = trt.Weights(np.ascontiguousarray(weight, dtype=dtype)) + conv_b = trt.Weights() + if bias is not None: + conv_b = trt.Weights(np.ascontiguousarray(bias, dtype=dtype)) + + conv = network.add_convolution_nd( + inp, + num_output_maps=out_channels, + kernel_shape=kernel_size, + kernel=conv_w, + bias=conv_b, + ) + conv.stride_nd = stride + conv.padding_nd = padding + conv.num_groups = groups + return conv.get_output(0) + + +def add_batch_norm( + network, + x, + gamma: np.ndarray, + beta: np.ndarray, + running_mean: np.ndarray, + running_var: np.ndarray, + eps: float, + *, + dtype=np.float32, +): + """Fold inference-time batch norm into a single per-channel scale+shift. + + y = (x - mean) / sqrt(var + eps) * gamma + beta + = x * scale + shift + Folding avoids emitting a normalization layer whose statistics are + constant at inference time. + """ + scale = (gamma / np.sqrt(running_var + eps)).astype(np.float32) + shift = (beta - running_mean * scale).astype(np.float32) + layer = network.add_scale( + x, + trt.ScaleMode.CHANNEL, + shift=trt.Weights(np.ascontiguousarray(shift, dtype=dtype)), + scale=trt.Weights(np.ascontiguousarray(scale, dtype=dtype)), + ) + return layer.get_output(0) + + +def add_relu(network, x): + return network.add_activation(x, trt.ActivationType.RELU).get_output(0) + + +def add_max_pool2d(network, x, kernel: int, stride: int, padding: int): + pool = network.add_pooling_nd(x, trt.PoolingType.MAX, (kernel, kernel)) + pool.stride_nd = (stride, stride) + pool.padding_nd = (padding, padding) + return pool.get_output(0) + + +def add_global_avg_pool(network, x, spatial: tuple[int, int]): + pool = network.add_pooling_nd(x, trt.PoolingType.AVERAGE, spatial) + pool.stride_nd = (1, 1) + return pool.get_output(0) + + +def add_sum(network, a, b): + return network.add_elementwise(a, b, trt.ElementWiseOperation.SUM).get_output(0) + + +def add_fc( + network, + x, + in_features: int, + out_features: int, + weight: np.ndarray, + bias: np.ndarray, + *, + dtype=np.float32, +): + """Final classifier: flatten the pooled feature map, then y = x @ W^T + b.""" + flat = network.add_shuffle(x) + flat.reshape_dims = (1, in_features) + flat_out = flat.get_output(0) + + # timm stores fc.weight as (out, in); TensorRT wants the (in, out) operand. + w = np.ascontiguousarray(weight.T, dtype=dtype) + w_const = network.add_constant((in_features, out_features), trt.Weights(w)).get_output(0) + mm = network.add_matrix_multiply( + flat_out, + trt.MatrixOperation.NONE, + w_const, + trt.MatrixOperation.NONE, + ).get_output(0) + + b = np.ascontiguousarray(bias.reshape(1, out_features), dtype=dtype) + b_const = network.add_constant((1, out_features), trt.Weights(b)).get_output(0) + return network.add_elementwise(mm, b_const, trt.ElementWiseOperation.SUM).get_output(0) diff --git a/python/tensorrt_model_connect/families/timm_repvgg/plugin.py b/python/tensorrt_model_connect/families/timm_repvgg/plugin.py new file mode 100644 index 0000000000..b244ac0956 --- /dev/null +++ b/python/tensorrt_model_connect/families/timm_repvgg/plugin.py @@ -0,0 +1,350 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""timm RepVGG image-classification family plugin. + +Supports timm RepVGG classifiers stored in HF Hub format. The initial target is: + timm/repvgg_a2.rvgg_in1k + +The builder constructs the classifier with TensorRT Network API calls rather +than routing through ONNX, matching the other timm families. + +The published checkpoints are in RepVGG's *training* form: every block keeps a +3x3 branch, a 1x1 branch, and, when the shape is unchanged, a batch-norm +identity branch. This loader performs the structural reparameterisation on the +host, folding all three into a single 3x3 convolution with bias, which is how +RepVGG is meant to run. The engine is therefore a plain convolution stack. + +The layout is recovered from the checkpoint, including the stride: a block +downsamples exactly when it has no identity branch, which is the only case where +its input and output shapes can differ. +""" + +from __future__ import annotations + +import re +import sys +from pathlib import Path + +import numpy as np +from tensorrt_model_connect import trt_compat + +from .model import model as graph_ops +from .weights import ( + WeightDict, + _has_tensor, + _load_tensor, + _open_safetensors, + _target_np_dtype, +) +from .config import ModelConfig + + +trt = trt_compat.get_trt() + +_BN_EPS = 1e-5 + + +def _pretrained_cfg(raw: dict) -> dict: + nested = raw.get("pretrained_cfg") + return nested if isinstance(nested, dict) else raw + + +def _resolve_config(raw: dict) -> dict: + pcfg = _pretrained_cfg(raw) + input_size = pcfg.get("input_size", [3, 224, 224]) + if isinstance(input_size, int): + image_h = image_w = int(input_size) + else: + image_h, image_w = int(input_size[-2]), int(input_size[-1]) + return { + "image_size_h": image_h, + "image_size_w": image_w, + "num_classes": int(raw.get("num_classes", pcfg.get("num_classes", 1000))), + "num_features": int(raw.get("num_features", 1408)), + "mean": [float(v) for v in pcfg.get("mean", [0.485, 0.456, 0.406])], + "std": [float(v) for v in pcfg.get("std", [0.229, 0.224, 0.225])], + "crop_pct": float(pcfg.get("crop_pct", 0.875)), + "interpolation": str(pcfg.get("interpolation", "bilinear")), + } + + +def _fold_conv_bn( + weight: np.ndarray, + gamma: np.ndarray, + beta: np.ndarray, + mean: np.ndarray, + var: np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: + """Fold an inference-time batch norm into its convolution.""" + scale = gamma / np.sqrt(var + _BN_EPS) + folded = weight * scale.reshape(-1, 1, 1, 1) + return folded.astype(np.float32), (beta - mean * scale).astype(np.float32) + + +def _identity_kernel( + channels: int, + gamma: np.ndarray, + beta: np.ndarray, + mean: np.ndarray, + var: np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: + """The identity branch as an equivalent 3x3 convolution. + + A batch norm applied straight to the input is a 1x1 identity convolution + scaled per channel, which pads into the centre of a 3x3 kernel. + """ + scale = gamma / np.sqrt(var + _BN_EPS) + kernel = np.zeros((channels, channels, 3, 3), dtype=np.float32) + for channel in range(channels): + kernel[channel, channel, 1, 1] = scale[channel] + return kernel, (beta - mean * scale).astype(np.float32) + + +def _discover_layout(readers) -> dict: + tensor_map = getattr(readers, "tensor_map", None) + if tensor_map is not None: + names = set(tensor_map) + else: + names = set() + for reader in readers: + names.update(reader.keys()) + + leaves: dict[tuple[int, int], set[str]] = {} + pattern = re.compile(r"^stages\.(\d+)\.(\d+)\.(.+)$") + for name in names: + match = pattern.match(name) + if match: + key = (int(match.group(1)), int(match.group(2))) + leaves.setdefault(key, set()).add(match.group(3).split(".")[0]) + if not leaves: + raise ValueError("Checkpoint has no stages.. entries") + + stages = sorted({stage for stage, _ in leaves}) + if stages != list(range(len(stages))): + raise ValueError("Stage indices are not contiguous") + + blocks = [] + for stage in stages: + indices = sorted(index for s, index in leaves if s == stage) + if indices != list(range(len(indices))): + raise ValueError(f"Stage {stage} block indices are not contiguous") + for index in indices: + present = leaves[(stage, index)] + if "conv_kxk" not in present: + raise ValueError(f"stages.{stage}.{index} has no conv_kxk branch") + has_identity = "identity" in present + blocks.append( + { + "prefix": f"stages.{stage}.{index}", + "has_identity": has_identity, + # Only a block without an identity branch may change shape. + "stride": 1 if has_identity else 2, + } + ) + return {"blocks": blocks, "num_stages": len(stages)} + + +class TimmRepvggPlugin: + name = "timm_repvgg" + runtime_strategy = "timm_repvgg_image_classification" + requires_tokenizer = False + + def matches(self, model_type: str) -> bool: + mt = (model_type or "").lower() + if mt == "timm_repvgg": + return True + return mt.startswith("repvgg") + + def _reparameterise( + self, readers, prefix: str, has_identity: bool, target_dtype + ) -> tuple[np.ndarray, np.ndarray]: + """Fuse the 3x3, 1x1 and identity branches into one 3x3 convolution.""" + + def bn(leaf: str): + return tuple( + _load_tensor(readers, f"{leaf}.{suffix}").astype(np.float32) + for suffix in ("weight", "bias", "running_mean", "running_var") + ) + + kxk = _load_tensor(readers, f"{prefix}.conv_kxk.conv.weight").astype(np.float32) + out_ch, in_per_group, kh, kw = kxk.shape + if (kh, kw) != (3, 3): + raise ValueError(f"{prefix}: expected a 3x3 branch, found {kh}x{kw}") + + gamma, beta, mean, var = bn(f"{prefix}.conv_kxk.bn") + fused_w, fused_b = _fold_conv_bn(kxk, gamma, beta, mean, var) + + one = _load_tensor(readers, f"{prefix}.conv_1x1.conv.weight").astype(np.float32) + if one.shape[1] != in_per_group: + raise ValueError( + f"{prefix}: the 1x1 and 3x3 branches disagree on grouping") + gamma, beta, mean, var = bn(f"{prefix}.conv_1x1.bn") + one_w, one_b = _fold_conv_bn(one, gamma, beta, mean, var) + # A 1x1 kernel is a 3x3 kernel with only its centre populated. + padded = np.zeros_like(fused_w) + padded[:, :, 1:2, 1:2] = one_w + fused_w = fused_w + padded + fused_b = fused_b + one_b + + if has_identity: + if in_per_group != out_ch: + raise ValueError( + f"{prefix}: an identity branch needs matching channel counts; " + f"grouped RepVGG variants are not supported") + gamma, beta, mean, var = bn(f"{prefix}.identity") + id_w, id_b = _identity_kernel(out_ch, gamma, beta, mean, var) + fused_w = fused_w + id_w + fused_b = fused_b + id_b + + return fused_w.astype(target_dtype), fused_b.astype(target_dtype) + + def load_weights( + self, + model_dir: str, + config: ModelConfig, + *, + precision: str = "fp32", + ) -> WeightDict: + readers = _open_safetensors(Path(model_dir)) + raw = config.raw + cfg = _resolve_config(raw) + layout = _discover_layout(readers) + cfg.update(layout) + raw["_timm_repvgg_config"] = cfg + target_dtype = _target_np_dtype(precision) + + weights = WeightDict() + + # The stem is the same multi-branch block without an identity path. + stem_w, stem_b = self._reparameterise(readers, "stem", False, target_dtype) + weights["stem.weight"] = stem_w + weights["stem.bias"] = stem_b + + for block in layout["blocks"]: + prefix = block["prefix"] + fused_w, fused_b = self._reparameterise( + readers, prefix, block["has_identity"], target_dtype) + weights[f"{prefix}.weight"] = fused_w + weights[f"{prefix}.bias"] = fused_b + + for key in ("head.fc.weight", "head.fc.bias"): + if not _has_tensor(readers, key): + raise KeyError(f"Tensor not found: {key}") + weights[key] = _load_tensor(readers, key).astype(target_dtype) + + return weights + + def build_engine( + self, + config: ModelConfig, + weights: WeightDict, + max_cache_length: int, + *, + precision: str = "fp32", + quant_ctx=None, + verbose: bool = False, + parallel_config=None, + ) -> bytes: + del max_cache_length + if quant_ctx is not None: + raise NotImplementedError("timm_repvgg does not support quantized builds yet") + if parallel_config is not None and getattr(parallel_config, "enabled", False): + raise NotImplementedError("timm_repvgg does not support tensor-parallel builds") + + if precision == "fp16": + work_np_dtype, work_trt_dtype = np.float16, trt.float16 + elif precision == "fp32": + work_np_dtype, work_trt_dtype = np.float32, trt.float32 + else: + raise ValueError(f"Unsupported timm_repvgg precision: {precision}") + + cfg = config.raw.get("_timm_repvgg_config") + if cfg is None: + raise RuntimeError( + "load_weights must run before build_engine to resolve the layout") + image_h = cfg["image_size_h"] + image_w = cfg["image_size_w"] + num_classes = cfg["num_classes"] + blocks = cfg["blocks"] + + total_stride = 2 + for block in blocks: + total_stride *= block["stride"] + if image_h % total_stride != 0 or image_w % total_stride != 0: + raise ValueError( + f"timm_repvgg input {image_h}x{image_w} must be divisible by {total_stride}") + feat_h, feat_w = image_h // total_stride, image_w // total_stride + + if verbose: + print( + "[trtmc build] timm_repvgg: " + f"image={image_h}x{image_w}, blocks={len(blocks)}, " + f"classes={num_classes}, precision={precision} " + "(branches fused into single 3x3 convolutions)", + file=sys.stderr, + ) + + logger = trt.Logger(trt.Logger.VERBOSE if verbose else trt.Logger.WARNING) + builder = trt.Builder(logger) + network = builder.create_network( + trt_compat.network_creation_flags(strongly_typed=True, explicit_batch=True)) + trt_config = builder.create_builder_config() + trt_config.avg_timing_iterations = 8 + trt_config.max_aux_streams = 0 + trt_config.set_flag(trt.BuilderFlag.DISABLE_TIMING_CACHE) + trt_config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 4 << 30) + + pixel_values = network.add_input( + "pixel_values", trt.float32, (1, 3, image_h, image_w)) + hidden = pixel_values + if hidden.dtype != work_trt_dtype: + hidden = network.add_cast(hidden, work_trt_dtype).get_output(0) + + stem_w = weights["stem.weight"] + hidden = graph_ops.add_conv2d( + network, hidden, stem_w, weights["stem.bias"], int(stem_w.shape[0]), (3, 3), + stride=(2, 2), padding=(1, 1), dtype=work_np_dtype) + hidden = graph_ops.add_relu(network, hidden) + + for block in blocks: + prefix = block["prefix"] + stride = block["stride"] + w = weights[f"{prefix}.weight"] + hidden = graph_ops.add_conv2d( + network, hidden, w, weights[f"{prefix}.bias"], int(w.shape[0]), (3, 3), + stride=(stride, stride), padding=(1, 1), dtype=work_np_dtype) + hidden = graph_ops.add_relu(network, hidden) + + hidden = graph_ops.add_global_avg_pool(network, hidden, (feat_h, feat_w)) + + head_w = weights["head.fc.weight"] + logits = graph_ops.add_fc( + network, hidden, int(head_w.shape[1]), num_classes, + head_w, weights["head.fc.bias"], dtype=work_np_dtype) + if logits.dtype != trt.float32: + logits = network.add_cast(logits, trt.float32).get_output(0) + logits.name = "logits" + network.mark_output(logits) + + plan = builder.build_serialized_network(network, trt_config) + if plan is None: + raise RuntimeError("TensorRT timm_repvgg engine build failed") + return bytes(plan) + + def get_bundle_config_overrides(self, config: ModelConfig) -> dict: + cfg = config.raw.get("_timm_repvgg_config") or _resolve_config(config.raw) + return { + "model_type": config.model_type, + "runtime_strategy": self.runtime_strategy, + "input_image_h": cfg["image_size_h"], + "input_image_w": cfg["image_size_w"], + "num_classes": cfg["num_classes"], + "image_mean": cfg["mean"], + "image_std": cfg["std"], + "crop_pct": cfg["crop_pct"], + "interpolation": cfg["interpolation"], + } + + +plugin = TimmRepvggPlugin() diff --git a/python/tensorrt_model_connect/families/timm_repvgg/python_profile_requirements/timm_repvgg_reference.lock.txt b/python/tensorrt_model_connect/families/timm_repvgg/python_profile_requirements/timm_repvgg_reference.lock.txt new file mode 100644 index 0000000000..260c4eb606 --- /dev/null +++ b/python/tensorrt_model_connect/families/timm_repvgg/python_profile_requirements/timm_repvgg_reference.lock.txt @@ -0,0 +1 @@ +timm==1.0.28 diff --git a/python/tensorrt_model_connect/families/timm_repvgg/python_profile_verify.py b/python/tensorrt_model_connect/families/timm_repvgg/python_profile_verify.py new file mode 100644 index 0000000000..03ac67651e --- /dev/null +++ b/python/tensorrt_model_connect/families/timm_repvgg/python_profile_verify.py @@ -0,0 +1,11 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from importlib.metadata import version + +import timm + +assert version("timm") == "1.0.28" +assert timm.__version__ == "1.0.28" +assert callable(timm.create_model) +print(f"timm={timm.__version__} create_model=ok") diff --git a/python/tensorrt_model_connect/families/timm_repvgg/weights/__init__.py b/python/tensorrt_model_connect/families/timm_repvgg/weights/__init__.py new file mode 100644 index 0000000000..e6d1daebdd --- /dev/null +++ b/python/tensorrt_model_connect/families/timm_repvgg/weights/__init__.py @@ -0,0 +1,199 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""1:1 port of standard_checkpoint_mapper.cpp + tensor_math.cpp to Python. + +Loads HF safetensors and maps keys to the flat weight dict expected by +standard_decoder_builder.py. All projections are transposed from HF +[out, in] layout to [in, out] for TRT matmul. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np + +# Register bfloat16 dtype with numpy (needed for safetensors without torch). +try: + import ml_dtypes # noqa: F401 +except ImportError: + pass + +from safetensors import safe_open + + +def _target_np_dtype(precision: str) -> np.dtype: + """Map precision string to numpy dtype for weight storage.""" + if precision in ("fp16", "bf16"): + return np.float16 + return np.float32 + + +def _transpose_2d(arr: np.ndarray, name: str, precision: str = "fp32") -> np.ndarray: + """Transpose [rows, cols] -> [cols, rows] in C-contiguous target dtype.""" + if arr.ndim != 2: + raise ValueError(f"Expected rank-2 tensor for transpose: {name}") + return np.ascontiguousarray(arr.T, dtype=_target_np_dtype(precision)) + + +class WeightDict(dict): + """A dict mapping logical weight names to flat float32 arrays. + + Keys follow the convention used by standard_decoder_builder.py: + - embedding: [vocab, hidden] + - layer.{i}.input_norm: [hidden] + - layer.{i}.w_q: [hidden, attention_size] + - layer.{i}.w_k: [hidden, kv_attention_size] + - layer.{i}.w_v: [hidden, kv_attention_size] + - layer.{i}.q_bias: [attention_size] (optional) + - layer.{i}.k_bias: [kv_attention_size] (optional) + - layer.{i}.v_bias: [kv_attention_size] (optional) + - layer.{i}.q_norm: [attention_size] (optional) + - layer.{i}.k_norm: [kv_attention_size] (optional) + - layer.{i}.w_o: [attention_size, hidden] + - layer.{i}.post_attn_norm: [hidden] + - layer.{i}.w_gate: [hidden, mlp_size] + - layer.{i}.w_up: [hidden, mlp_size] + - layer.{i}.w_down: [mlp_size, hidden] + - final_norm: [hidden] + - w_out: [hidden, vocab] + """ + + +# --------------------------------------------------------------------------- +# Safetensors I/O helpers +# --------------------------------------------------------------------------- + + +def _detect_framework() -> str: + """Use 'torch' if available (handles BF16 natively), else 'numpy'.""" + try: + import torch # noqa: F401 + + return "torch" + except ImportError: + return "numpy" + + +class _TorchBinReader: + """Adapter that wraps a pytorch .bin state dict with the safetensors reader + interface (keys() / get_tensor()).""" + + def __init__(self, path: Path): + import torch + + self._state = torch.load(str(path), map_location="cpu", weights_only=True) + + def keys(self) -> list[str]: + return list(self._state.keys()) + + def get_tensor(self, name: str): + return self._state[name] + + +class _ReaderCollection(list): + """Reader list with a cached tensor-name -> reader lookup table.""" + + def __init__(self, readers: list, *, tensor_map: dict[str, object] | None = None): + super().__init__(readers) + if tensor_map is None: + tensor_map = {} + for reader in readers: + for key in reader.keys(): + tensor_map[key] = reader + self.tensor_map = tensor_map + + +def _open_safetensors(model_dir: Path) -> list: + """Open all safetensor shards (or pytorch .bin) in a model directory.""" + fw = _detect_framework() + single = model_dir / "model.safetensors" + if single.exists(): + return _ReaderCollection([safe_open(str(single), framework=fw)]) + + index_path = model_dir / "model.safetensors.index.json" + if index_path.exists(): + import json + + index = json.loads(index_path.read_text()) + weight_map = index.get("weight_map", {}) + shard_files = sorted(set(weight_map.values())) + readers_by_file = { + shard: safe_open(str(model_dir / shard), framework=fw) for shard in shard_files + } + tensor_map = {name: readers_by_file[shard] for name, shard in weight_map.items()} + return _ReaderCollection( + [readers_by_file[shard] for shard in shard_files], + tensor_map=tensor_map, + ) + + # Diffusers format: diffusion_pytorch_model.safetensors + diff_single = model_dir / "diffusion_pytorch_model.safetensors" + if diff_single.exists(): + return _ReaderCollection([safe_open(str(diff_single), framework=fw)]) + + diff_index = model_dir / "diffusion_pytorch_model.safetensors.index.json" + if diff_index.exists(): + import json + + index = json.loads(diff_index.read_text()) + weight_map = index.get("weight_map", {}) + shard_files = sorted(set(weight_map.values())) + readers_by_file = { + shard: safe_open(str(model_dir / shard), framework=fw) for shard in shard_files + } + tensor_map = {name: readers_by_file[shard] for name, shard in weight_map.items()} + return _ReaderCollection( + [readers_by_file[shard] for shard in shard_files], + tensor_map=tensor_map, + ) + + # Fallback: pytorch_model.bin (older HF models) + bin_single = model_dir / "pytorch_model.bin" + if bin_single.exists(): + return _ReaderCollection([_TorchBinReader(bin_single)]) + + raise FileNotFoundError( + f"No model.safetensors, index.json, or pytorch_model.bin in {model_dir}" + ) + + +def _has_tensor(readers: list, name: str) -> bool: + tensor_map = getattr(readers, "tensor_map", None) + if tensor_map is not None: + return name in tensor_map + for r in readers: + if name in r.keys(): + return True + return False + + +def _to_numpy_fp32(t) -> np.ndarray: + """Convert a safetensors/torch tensor to numpy float32 with minimal copies.""" + if hasattr(t, "numpy"): + dtype = getattr(t, "dtype", None) + if str(dtype) == "torch.float32": + return t.numpy() + return t.float().numpy() + + dtype_str = str(t.dtype) + if t.dtype == np.uint16 or dtype_str == "bfloat16": + t = t.view(np.uint16).astype(np.uint32) << 16 + return t.view(np.float32) + if dtype_str == "float16": + return t.astype(np.float32) + return np.asarray(t, dtype=np.float32) + + +def _load_tensor(readers: list, name: str) -> np.ndarray: + tensor_map = getattr(readers, "tensor_map", None) + if tensor_map is not None: + reader = tensor_map.get(name) + if reader is None: + raise KeyError(f"Tensor not found: {name}") + return _to_numpy_fp32(reader.get_tensor(name)) + for r in readers: + if name in r.keys(): + return _to_numpy_fp32(r.get_tensor(name)) + raise KeyError(f"Tensor not found: {name}") diff --git a/src/runtime/models/timm_repvgg/MODEL.toml b/src/runtime/models/timm_repvgg/MODEL.toml new file mode 100644 index 0000000000..8b5c2dd728 --- /dev/null +++ b/src/runtime/models/timm_repvgg/MODEL.toml @@ -0,0 +1,10 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +id = "timm_repvgg" +runtime_library = "libtrtmc_model_timm_repvgg.so" +runtime_plugins = ["plugin.cpp|register_timm_repvgg_plugin"] +runtime_strategies = ["timm_repvgg_image_classification"] +runtime_tests = [ + "test_timm_repvgg_image_preprocess_seam|test_timm_repvgg_image_preprocess_seam.cpp|trtmc_model_timm_repvgg|_|_", +] diff --git a/src/runtime/models/timm_repvgg/image_preprocess_seam.cpp b/src/runtime/models/timm_repvgg/image_preprocess_seam.cpp new file mode 100644 index 0000000000..799434b943 --- /dev/null +++ b/src/runtime/models/timm_repvgg/image_preprocess_seam.cpp @@ -0,0 +1,130 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/models/timm_repvgg/image_preprocess_seam.h" + +#include "stb_image_resize2.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace trtmc { + +namespace { + +stbir_filter resolve_timm_repvgg_resize_filter(const std::string& interpolation) { + if (interpolation == "bilinear") + return STBIR_FILTER_TRIANGLE; + if (interpolation == "bicubic") + return STBIR_FILTER_CATMULLROM; + throw std::invalid_argument("Unsupported timm RepVGG interpolation: " + interpolation); +} + +void validate_timm_repvgg_preprocess_config(const TimmRepvggPreprocessConfig& config) { + if (config.input_image_h <= 0 || config.input_image_w <= 0) { + throw std::invalid_argument("timm RepVGG input dimensions must be positive"); + } + if (config.crop_pct <= 0.0F || config.crop_pct > 1.0F) { + throw std::invalid_argument("timm RepVGG crop_pct must be in (0, 1]"); + } + if (config.image_mean.size() != 3 || config.image_std.size() != 3) { + throw std::invalid_argument("timm RepVGG image mean/std must contain three channels"); + } + for (float value : config.image_std) { + if (value == 0.0F) + throw std::invalid_argument("timm RepVGG image std must be non-zero"); + } +} + +int32_t torchvision_center_crop_offset(int32_t resized, int32_t target) { + const int32_t difference = resized - target; + const int32_t half = difference / 2; + // torchvision uses Python round(), whose .5 ties round to the nearest even integer. + return (difference % 2 != 0 && half % 2 != 0) ? half + 1 : half; +} + +} // namespace + +TimmRepvggResizeShape compute_timm_repvgg_resize_shape(int32_t image_height, int32_t image_width, + const TimmRepvggPreprocessConfig& config) { + if (image_height <= 0 || image_width <= 0) { + throw std::invalid_argument("timm RepVGG source dimensions must be positive"); + } + validate_timm_repvgg_preprocess_config(config); + + if (config.input_image_h == config.input_image_w) { + // timm's center-crop eval transform passes floor(input_size / crop_pct) as a scalar + // torchvision Resize size. A scalar fixes the shorter edge and floors the aspect-ratio + // calculation for the longer edge. + const int32_t resized_short = static_cast( + std::floor(static_cast(config.input_image_h) / config.crop_pct)); + if (image_height <= image_width) { + return {resized_short, static_cast(static_cast(resized_short) * + image_width / image_height)}; + } + return { + static_cast(static_cast(resized_short) * image_height / image_width), + resized_short}; + } + + const float required_scale = + std::max(static_cast(config.input_image_h) / static_cast(image_height), + static_cast(config.input_image_w) / static_cast(image_width)); + const float resize_scale = required_scale / config.crop_pct; + return { + std::max(config.input_image_h, + static_cast(std::floor(static_cast(image_height) * resize_scale))), + std::max(config.input_image_w, + static_cast(std::floor(static_cast(image_width) * resize_scale))), + }; +} + +std::vector preprocess_timm_repvgg_image(const float* image_pixels, int32_t image_height, + int32_t image_width, + const TimmRepvggPreprocessConfig& config) { + if (image_pixels == nullptr || image_height <= 0 || image_width <= 0) { + throw std::invalid_argument("timm RepVGG source image must be non-empty"); + } + validate_timm_repvgg_preprocess_config(config); + + const auto resize_shape = compute_timm_repvgg_resize_shape(image_height, image_width, config); + const int32_t resized_h = resize_shape.height; + const int32_t resized_w = resize_shape.width; + + std::vector resized(static_cast(resized_h) * resized_w * 3U); + if (stbir_resize(image_pixels, image_width, image_height, + image_width * 3 * static_cast(sizeof(float)), resized.data(), + resized_w, resized_h, resized_w * 3 * static_cast(sizeof(float)), + STBIR_RGB, STBIR_TYPE_FLOAT, STBIR_EDGE_CLAMP, + resolve_timm_repvgg_resize_filter(config.interpolation)) == nullptr) { + throw std::runtime_error("Failed to resize timm RepVGG input image"); + } + + const int32_t crop_y = torchvision_center_crop_offset(resized_h, config.input_image_h); + const int32_t crop_x = torchvision_center_crop_offset(resized_w, config.input_image_w); + const auto output_plane = static_cast(config.input_image_h) * config.input_image_w; + std::vector pixel_values(3U * output_plane); + for (int32_t y = 0; y < config.input_image_h; ++y) { + for (int32_t x = 0; x < config.input_image_w; ++x) { + const auto src_idx = + static_cast((((crop_y + y) * resized_w + crop_x + x) * 3)); + for (int32_t c = 0; c < 3; ++c) { + const auto channel = static_cast(c); + pixel_values[channel * output_plane + + static_cast(y) * config.input_image_w + x] = + (resized[src_idx + channel] - config.image_mean[channel]) / + config.image_std[channel]; + } + } + } + return pixel_values; +} + +} // namespace trtmc diff --git a/src/runtime/models/timm_repvgg/image_preprocess_seam.h b/src/runtime/models/timm_repvgg/image_preprocess_seam.h new file mode 100644 index 0000000000..165f8ce939 --- /dev/null +++ b/src/runtime/models/timm_repvgg/image_preprocess_seam.h @@ -0,0 +1,35 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +namespace trtmc { + +struct TimmRepvggPreprocessConfig { + int32_t input_image_h{224}; + int32_t input_image_w{224}; + std::vector image_mean{0.5F, 0.5F, 0.5F}; + std::vector image_std{0.5F, 0.5F, 0.5F}; + float crop_pct{0.9F}; + std::string interpolation{"bicubic"}; +}; + +struct TimmRepvggResizeShape { + int32_t height{0}; + int32_t width{0}; +}; + +TimmRepvggResizeShape compute_timm_repvgg_resize_shape(int32_t image_height, int32_t image_width, + const TimmRepvggPreprocessConfig& config); + +std::vector preprocess_timm_repvgg_image(const float* image_pixels, int32_t image_height, + int32_t image_width, + const TimmRepvggPreprocessConfig& config); + +} // namespace trtmc diff --git a/src/runtime/models/timm_repvgg/pipeline.cpp b/src/runtime/models/timm_repvgg/pipeline.cpp new file mode 100644 index 0000000000..4782defbc0 --- /dev/null +++ b/src/runtime/models/timm_repvgg/pipeline.cpp @@ -0,0 +1,67 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/models/timm_repvgg/pipeline.h" + +#include +#include +#include +#include + +namespace trtmc { + +namespace { + +const Tensor* find_logits_output(const TensorMap& outputs) { + for (const auto& [name, tensor] : outputs) { + if (name.find("logits") != std::string::npos || outputs.size() == 1) + return &tensor; + } + return nullptr; +} + +} // namespace + +TimmRepvggImageClassificationPipeline::TimmRepvggImageClassificationPipeline( + std::unique_ptr model, TimmRepvggPreprocessConfig preprocess_config, + std::string model_id_str) + : model_(std::move(model)), preprocess_config_(std::move(preprocess_config)), + model_id_(std::move(model_id_str)) { + if (!model_ || !model_->ok()) + throw std::runtime_error("TimmRepvggImageClassificationPipeline: invalid model"); +} + +ClassificationResult TimmRepvggImageClassificationPipeline::classify(const float* pixels, + int32_t height, + int32_t width) { + auto pixel_values = preprocess_timm_repvgg_image(pixels, height, width, preprocess_config_); + + Tensor img_t; + img_t.data = pixel_values.data(); + img_t.shape = {1, 3, preprocess_config_.input_image_h, preprocess_config_.input_image_w}; + img_t.dtype = DType::kFloat32; + + auto outputs = model_->forward({{"pixel_values", img_t}}); + ClassificationResult result; + + const Tensor* logits_tensor = find_logits_output(outputs); + if (!logits_tensor) + return result; + + const auto n = logits_tensor->numel(); + if (n <= 0) + return result; + + result.logits.resize(static_cast(n)); + std::memcpy(result.logits.data(), logits_tensor->data, + static_cast(n) * sizeof(float)); + + auto best = std::max_element(result.logits.begin(), result.logits.end()); + result.top_class = static_cast(std::distance(result.logits.begin(), best)); + result.top_score = (best == result.logits.end()) ? 0.0F : *best; + return result; +} + +} // namespace trtmc diff --git a/src/runtime/models/timm_repvgg/pipeline.h b/src/runtime/models/timm_repvgg/pipeline.h new file mode 100644 index 0000000000..25f56b5c53 --- /dev/null +++ b/src/runtime/models/timm_repvgg/pipeline.h @@ -0,0 +1,34 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "runtime/models/timm_repvgg/image_preprocess_seam.h" +#include "trtmc/pipeline.h" +#include "trtmc/runtime/trt_module.h" + +#include +#include + +namespace trtmc { + +class TimmRepvggImageClassificationPipeline final : public IPipeline { + public: + explicit TimmRepvggImageClassificationPipeline( + std::unique_ptr model, TimmRepvggPreprocessConfig preprocess_config = {}, + std::string model_id_str = ""); + + ClassificationResult classify(const float* pixels, int32_t height, int32_t width) override; + + const char* model_id() const override { return model_id_.c_str(); } + const char* pipeline_type() const override { return "TimmRepvggImageClassificationPipeline"; } + + private: + std::unique_ptr model_; + TimmRepvggPreprocessConfig preprocess_config_; + std::string model_id_; +}; + +} // namespace trtmc diff --git a/src/runtime/models/timm_repvgg/plugin.cpp b/src/runtime/models/timm_repvgg/plugin.cpp new file mode 100644 index 0000000000..9c911b2bf8 --- /dev/null +++ b/src/runtime/models/timm_repvgg/plugin.cpp @@ -0,0 +1,85 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +// TimmRepvggPlugin: handles timm RepVGG image-classification bundles. + +#include "plugin_helpers.h" +#include "runtime/models/timm_repvgg/pipeline.h" +#include "trtmc/runtime/distributed_runtime.h" +#include "trtmc/runtime/pipeline_registry.h" +#include "utils/json_helpers.h" + +#include +#include +#include + +namespace trtmc { + +namespace { + +struct TensorParallelRuntimeConfig { + bool enabled{false}; + int32_t tp_size{1}; +}; + +TensorParallelRuntimeConfig parse_tensor_parallel_runtime_config(const std::string& config_json) { + TensorParallelRuntimeConfig cfg; + cfg.tp_size = extract_json_int(config_json, "tensor_parallel_size", 1); + const auto mode = extract_json_string(config_json, "tensor_parallel_mode", "single"); + cfg.enabled = (mode == "tensor_parallel" && cfg.tp_size > 1); + return cfg; +} + +std::string tp_engine_section_name(int32_t rank) { + return "engine_plan_tp_rank" + std::to_string(rank); +} + +TimmRepvggPreprocessConfig make_timm_repvgg_preprocess_config(const std::string& json) { + TimmRepvggPreprocessConfig cfg; + cfg.input_image_h = extract_json_int(json, "input_image_h", cfg.input_image_h); + cfg.input_image_w = extract_json_int(json, "input_image_w", cfg.input_image_w); + cfg.crop_pct = extract_json_float(json, "crop_pct", cfg.crop_pct); + cfg.interpolation = extract_json_string(json, "interpolation", cfg.interpolation); + + auto mean = extract_json_float_array(json, "image_mean", 3); + if (mean.size() == 3) + cfg.image_mean = std::move(mean); + auto stdv = extract_json_float_array(json, "image_std", 3); + if (stdv.size() == 3) + cfg.image_std = std::move(stdv); + return cfg; +} + +} // namespace + +class TimmRepvggPlugin final : public IPipelinePlugin { + public: + std::unique_ptr create(const PipelineContext& ctx) override { + ModuleCreateOptions opts; + opts.runtime_cache_path = ctx.runtime_cache_path.c_str(); + opts.cuda_graphs = ctx.cuda_graphs; + + const auto tp_config = parse_tensor_parallel_runtime_config(ctx.config_json); + DistributedRuntimeGroup tp_group; + std::string engine_section = "engine_plan"; + if (tp_config.enabled) { + tp_group = initialize_tensor_parallel_group(tp_config.tp_size); + opts.distributed_communicator = tp_group.communicator; + opts.distributed_owner = tp_group.owner; + engine_section = tp_engine_section_name(tp_group.rank); + } + + auto loaded = load_trt_module_from_plan( + ctx.backend, find_section(ctx.bundle, engine_section), "engine_plan", opts); + return std::make_unique( + std::move(loaded.module), make_timm_repvgg_preprocess_config(ctx.config_json), + ctx.bundle.info.model_id); + } +}; + +REGISTER_PIPELINE_PLUGIN_WITH_MANIFEST(register_timm_repvgg_plugin, TimmRepvggPlugin, + "timm_repvgg_image_classification"); + +} // namespace trtmc diff --git a/src/runtime/models/timm_repvgg/plugin_helpers.cpp b/src/runtime/models/timm_repvgg/plugin_helpers.cpp new file mode 100644 index 0000000000..2effdc9709 --- /dev/null +++ b/src/runtime/models/timm_repvgg/plugin_helpers.cpp @@ -0,0 +1,476 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "plugin_helpers.h" + +#include "trtmc/runtime/trt_backend.h" +#include "utils/json_helpers.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if TRTMC_HAS_TVM_FFI +#include "plugins/tvm_ffi_module_loader.h" +#endif + +namespace trtmc { + +namespace { + +using SteadyClock = std::chrono::steady_clock; + +double elapsed_ms(SteadyClock::time_point start, SteadyClock::time_point end) { + return std::chrono::duration(end - start).count(); +} + +class SpecialFrameTokenizer final : public ITokenizer { + public: + SpecialFrameTokenizer(std::shared_ptr inner, std::vector prefix, + std::vector suffix) + : mInner(std::move(inner)), mPrefix(std::move(prefix)), mSuffix(std::move(suffix)) {} + + std::vector encode(const std::string& text) const override { + auto ids = mInner->encode(text); + std::vector framed; + framed.reserve(mPrefix.size() + ids.size() + mSuffix.size()); + framed.insert(framed.end(), mPrefix.begin(), mPrefix.end()); + framed.insert(framed.end(), ids.begin(), ids.end()); + framed.insert(framed.end(), mSuffix.begin(), mSuffix.end()); + return framed; + } + + std::string decode(const std::vector& ids) const override { + return mInner->decode(ids); + } + + int32_t id_for_token(std::string_view token) const override { + return mInner->id_for_token(token); + } + + std::string token_for_id(int32_t id) const override { return mInner->token_for_id(id); } + + private: + std::shared_ptr mInner; + std::vector mPrefix; + std::vector mSuffix; +}; + +struct TokenizerSpecialFrame { + bool present{false}; + std::vector prefix; + std::vector suffix; +}; + +using TokenizerFactory = std::unique_ptr (*)(const char*, std::size_t, bool); + +} // namespace + +void log_trt_load_timing(const char* label, double load_deserialize_ms, std::size_t plan_bytes) { + std::ostringstream line; + line << std::fixed << std::setprecision(6) << "[trtmc.load_timing] label=\"" + << (label ? label : "engine") << "\" load_deserialize_ms=" << load_deserialize_ms + << " plan_bytes=" << plan_bytes; + std::cerr << line.str() << '\n'; +} + +// Tokenizer helpers. + +bool detect_add_special_tokens(const BundleFile& bundle) { + if (bundle.info.tokenizer_add_special_tokens_present) + return bundle.info.tokenizer_add_special_tokens; + + auto* config_data = find_section(bundle, "config.json"); + if (!config_data) + return true; + std::string cfg_text(config_data->begin(), config_data->end()); + auto pos = cfg_text.find("\"tokenizer_add_special_tokens\""); + if (pos == std::string::npos) + return true; + auto val_pos = cfg_text.find(':', pos); + if (val_pos == std::string::npos) + return true; + auto value_pos = cfg_text.find_first_not_of(" \t\r\n", val_pos + 1); + if (value_pos == std::string::npos) + return true; + if (cfg_text.compare(value_pos, 5, "false") == 0 || cfg_text[value_pos] == '0') + return false; + if (cfg_text.compare(value_pos, 4, "true") == 0 || cfg_text[value_pos] == '1') + return true; + return true; +} + +namespace { + +TokenizerSpecialFrame detect_tokenizer_special_frame(const BundleFile& bundle) { + TokenizerSpecialFrame frame; + auto* config_data = find_section(bundle, "config.json"); + if (!config_data) + return frame; + std::string cfg_text(config_data->begin(), config_data->end()); + const bool has_prefix = cfg_text.find("\"tokenizer_special_prefix_ids\"") != std::string::npos; + const bool has_suffix = cfg_text.find("\"tokenizer_special_suffix_ids\"") != std::string::npos; + if (!has_prefix && !has_suffix) + return frame; + + frame.present = true; + frame.prefix = extract_json_int_array(cfg_text, "tokenizer_special_prefix_ids"); + frame.suffix = extract_json_int_array(cfg_text, "tokenizer_special_suffix_ids"); + return frame; +} + +std::shared_ptr apply_tokenizer_special_frame(std::unique_ptr tokenizer, + const TokenizerSpecialFrame& frame) { + if (!tokenizer) + return nullptr; + std::shared_ptr shared(std::move(tokenizer)); + if (!frame.present || (frame.prefix.empty() && frame.suffix.empty())) + return shared; + return std::make_shared(std::move(shared), frame.prefix, frame.suffix); +} + +TokenizerSpecialFrame detect_requested_tokenizer_special_frame(const BundleFile& bundle, + bool add_special_tokens) { + if (!add_special_tokens) + return TokenizerSpecialFrame{}; + return detect_tokenizer_special_frame(bundle); +} + +std::shared_ptr try_create_native_tokenizer_kind(TokenizerFactory factory, + const char* data, std::size_t size, + bool add_special_tokens, + const TokenizerSpecialFrame& frame, + const char* label) { + try { + auto tok = factory(data, size, add_special_tokens); + if (!tok) + return nullptr; + std::cerr << "[trtmc] Using native " << label << " tokenizer" << std::endl; + return apply_tokenizer_special_frame(std::move(tok), frame); + } catch (...) { + return nullptr; + } +} + +} // namespace + +bool is_bpe_tokenizer_json(const BundleFile& bundle) { + auto* tok_data = find_section(bundle, "tokenizer.json"); + if (!tok_data || tok_data->empty()) + return false; + // Quick string search — avoid full JSON parse just for type detection + std::string_view json(tok_data->data(), tok_data->size()); + return json.find("\"type\":\"BPE\"") != std::string_view::npos || + json.find("\"type\": \"BPE\"") != std::string_view::npos; +} + +std::shared_ptr try_create_native_bpe(const BundleFile& bundle, bool add_special, + bool throw_on_failure) { + auto* tok_data = find_section(bundle, "tokenizer.json"); + if (!tok_data || tok_data->empty()) + return nullptr; + try { + auto tok = CreateBpeTokenizer(tok_data->data(), tok_data->size(), add_special); + if (tok) { + std::cerr << "[trtmc] Using native BPE tokenizer" << std::endl; + } + return tok; + } catch (const std::exception& e) { + // "Not a BPE tokenizer" -> non-BPE model (WordPiece, Unigram), allow fallback + std::string msg = e.what(); + bool is_non_bpe = msg.find("Not a BPE") != std::string::npos; + + if (throw_on_failure || (!is_non_bpe && is_bpe_tokenizer_json(bundle))) { + // BPE model but native failed -> error, no silent fallback + throw std::runtime_error(std::string("Native BPE tokenizer failed for BPE model: ") + + e.what()); + } + std::cerr << "[trtmc] Native BPE unavailable (" << e.what() + << "), falling back to HF Python" << std::endl; + } + return nullptr; +} + +std::shared_ptr try_create_native_tokenizer(const BundleFile& bundle, + bool add_special_tokens) { + auto* tok_data = find_section(bundle, "tokenizer.json"); + if (!tok_data || tok_data->empty()) + return nullptr; + + const char* data = tok_data->data(); + std::size_t size = tok_data->size(); + const auto special_frame = detect_requested_tokenizer_special_frame(bundle, add_special_tokens); + const bool native_add_special = !special_frame.present && add_special_tokens; + + if (auto tokenizer = try_create_native_tokenizer_kind(CreateBpeTokenizer, data, size, + native_add_special, special_frame, "BPE")) + return tokenizer; + + if (auto tokenizer = try_create_native_tokenizer_kind( + CreateWordPieceTokenizer, data, size, native_add_special, special_frame, "WordPiece")) + return tokenizer; + + return try_create_native_tokenizer_kind(CreateUnigramTokenizer, data, size, native_add_special, + special_frame, "Unigram"); +} + +std::shared_ptr create_tokenizer_from_bundle(const BundleFile& bundle) { + bool add_special = detect_add_special_tokens(bundle); + return try_create_native_tokenizer(bundle, add_special); +} + +// TRT module loading (delegated to IBackend). + +LoadedModule load_trt_module_from_plan(IBackend* backend, const std::vector* plan, + const char* label, const ModuleCreateOptions& options) { + if (!plan || plan->empty()) + throw std::runtime_error(std::string("Bundle missing ") + label); + if (!backend) + throw std::runtime_error("No backend loaded"); + + LoadedModule result; + const auto t0 = SteadyClock::now(); + result.module = backend->create_module(plan->data(), plan->size(), options); + const auto t1 = SteadyClock::now(); + log_trt_load_timing(label, elapsed_ms(t0, t1), plan->size()); + if (!result.module || !result.module->ok()) + throw std::runtime_error(std::string("Failed to create ITrtModule for ") + label); + result.module->set_timing_label(label ? label : "engine"); + return result; +} + +LoadedModule try_load_trt_module_from_plan(IBackend* backend, const std::vector* plan, + const char* label, const ModuleCreateOptions& options) { + if (!plan || plan->empty()) + return LoadedModule{}; + try { + return load_trt_module_from_plan(backend, plan, label, options); + } catch (...) { + std::cerr << "[trtmc] WARNING: failed to load optional engine: " << label << std::endl; + return LoadedModule{}; + } +} + +std::unique_ptr extract_optional_module(IBackend* backend, + const std::vector* plan, + const char* label, + const ModuleCreateOptions& options) { + auto loaded = try_load_trt_module_from_plan(backend, plan, label, options); + if (loaded.module && loaded.module->ok()) + return std::move(loaded.module); + return nullptr; +} + +// Dual-profile module loading (delegated to IBackend). + +DualProfileModules load_dual_profile_modules(IBackend* backend, const std::vector* plan, + const char* label, + const ModuleCreateOptions& options) { + if (!plan || plan->empty()) + throw std::runtime_error(std::string("Bundle missing ") + label); + if (!backend) + throw std::runtime_error("No backend loaded"); + + const auto t0 = SteadyClock::now(); + auto pair = backend->create_dual_profile_modules(plan->data(), plan->size(), options); + const auto t1 = SteadyClock::now(); + log_trt_load_timing(label, elapsed_ms(t0, t1), plan->size()); + if (!pair.decode || !pair.decode->ok()) + throw std::runtime_error(std::string("Failed to create dual-profile modules for ") + label); + + DualProfileModules out; + out.prefill = std::move(pair.prefill); + out.decode = std::move(pair.decode); + if (out.prefill) + out.prefill->set_timing_label(std::string(label ? label : "engine") + ":prefill"); + if (out.decode) + out.decode->set_timing_label(std::string(label ? label : "engine") + ":decode"); + return out; +} + +// Config helpers. + +int32_t compute_kv_dim(const BaseConfig& cfg) { + int32_t hd = (cfg.head_dim > 0) ? cfg.head_dim + : ((cfg.num_heads > 0) ? cfg.hidden_size / cfg.num_heads : 128); + int32_t kv_heads = (cfg.num_kv_heads > 0) ? cfg.num_kv_heads : cfg.num_heads; + return kv_heads * hd; +} + +DType cache_dtype_from_precision(const std::string& precision) { + if (precision == "fp16") + return DType::kFloat16; + if (precision == "bf16") + return DType::kBFloat16; + return DType::kFloat32; +} + +// Section data conversion. + +std::vector section_to_floats(const std::vector* sec) { + if (!sec || sec->empty()) + return {}; + std::size_t count = sec->size() / sizeof(float); + std::vector out(count); + std::memcpy(out.data(), sec->data(), count * sizeof(float)); + return out; +} + +std::vector section_to_int32s(const std::vector* sec) { + if (!sec || sec->empty()) + return {}; + std::size_t count = sec->size() / sizeof(int32_t); + std::vector out(count); + std::memcpy(out.data(), sec->data(), count * sizeof(int32_t)); + return out; +} + +bool has_section_data(const std::vector* d) { + return d && !d->empty(); +} + +MelFilterbank load_mel_filterbank(const BundleFile& bundle) { + MelFilterbank fb; + const auto* data = find_section(bundle, "mel_filterbank"); + if (data == nullptr || data->empty()) + return fb; + + // Format: [n_freq_bins(int32), n_mel_bins(int32), float32 data...] + if (data->size() < 2 * sizeof(int32_t)) + return fb; + + int32_t header[2] = {0, 0}; + std::memcpy(header, data->data(), sizeof(header)); + fb.n_freq_bins = header[0]; + fb.n_mel_bins = header[1]; + + if (fb.n_freq_bins <= 0 || fb.n_mel_bins <= 0) + return fb; + + const auto expected_data_size = static_cast(fb.n_freq_bins) * + static_cast(fb.n_mel_bins) * sizeof(float); + const auto payload_offset = 2 * sizeof(int32_t); + if (data->size() < payload_offset + expected_data_size) { + fb.n_freq_bins = 0; + fb.n_mel_bins = 0; + return fb; + } + + fb.data.resize(static_cast(fb.n_freq_bins) * fb.n_mel_bins); + std::memcpy(fb.data.data(), data->data() + payload_offset, expected_data_size); + return fb; +} + +std::unique_ptr create_clip_tokenizer_from_bundle(const BundleFile& bundle) { + auto* tok_data = find_section(bundle, "clip_tokenizer.json"); + if (!tok_data || tok_data->empty()) + return nullptr; + try { + auto tok = + CreateBpeTokenizer(tok_data->data(), tok_data->size(), /*add_special_tokens=*/true); + if (tok) + std::cerr << "[trtmc] Using native BPE CLIP tokenizer" << std::endl; + return tok; + } catch (const std::exception& e) { + std::cerr << "[trtmc] WARNING: CLIP tokenizer failed: " << e.what() << std::endl; + } + return nullptr; +} + +// ─── FFI kernel loading ─── + +#if TRTMC_HAS_TVM_FFI + +namespace { + +// Write a bundle section to a temporary .so file, returning the path. +std::string write_kernel_so_to_temp(const std::string& global_name, const char* data, + std::size_t size) { + std::string safe_name = global_name; + for (auto& c : safe_name) { + if (c == '.') + c = '_'; + } + std::string tmp_path = "/tmp/trtmc_kernel_" + safe_name + ".so"; + std::ofstream ofs(tmp_path, std::ios::binary); + ofs.write(data, static_cast(size)); + return tmp_path; +} + +// Load a single kernel entry from the manifest and register it via TVM-FFI. +void load_single_kernel(const BundleFile& bundle, const std::string& obj) { + std::string global_name = extract_json_string(obj, "global_name", ""); + std::string func_name = extract_json_string(obj, "func_name", "run"); + std::string section_name = extract_json_string(obj, "section", ""); + + if (global_name.empty() || section_name.empty()) + return; + + const auto* so_sec = find_section(bundle, section_name); + if (!so_sec || so_sec->empty()) { + std::cerr << "[ffi] Kernel .so section not found: " << section_name << '\n'; + return; + } + + std::string tmp_path = write_kernel_so_to_temp(global_name, so_sec->data(), so_sec->size()); + if (load_tvm_ffi_module_func(tmp_path, func_name, global_name)) { + std::cerr << "[ffi] Loaded kernel: " << global_name << '\n'; + } else { + std::cerr << "[ffi] Failed to load kernel: " << global_name << " from " << section_name + << '\n'; + } +} + +// Find the "kernels" JSON array bounds within the manifest string. +// Returns {start_after_bracket, closing_bracket} or {npos, npos}. +std::pair find_kernels_array_bounds(const std::string& s) { + auto pos = s.find("\"kernels\""); + if (pos == std::string::npos) + return {std::string::npos, std::string::npos}; + auto arr_start = s.find('[', pos); + if (arr_start == std::string::npos) + return {std::string::npos, std::string::npos}; + auto arr_end = s.find(']', arr_start); + return {arr_start + 1, arr_end}; +} + +} // namespace + +#endif // TRTMC_HAS_TVM_FFI + +void load_ffi_kernels_from_bundle(const BundleFile& bundle) { +#if TRTMC_HAS_TVM_FFI + const auto* manifest_sec = find_section(bundle, "kernel_manifest.json"); + if (!manifest_sec) + return; + + std::string manifest_str(manifest_sec->begin(), manifest_sec->end()); + auto [cur, arr_end] = find_kernels_array_bounds(manifest_str); + if (cur == std::string::npos || arr_end == std::string::npos) + return; + + while (cur < arr_end) { + auto obj_start = manifest_str.find('{', cur); + if (obj_start == std::string::npos || obj_start >= arr_end) + break; + auto obj_end = manifest_str.find('}', obj_start); + if (obj_end == std::string::npos) + break; + + load_single_kernel(bundle, manifest_str.substr(obj_start, obj_end - obj_start + 1)); + cur = obj_end + 1; + } +#else + (void)bundle; +#endif +} + +} // namespace trtmc diff --git a/src/runtime/models/timm_repvgg/plugin_helpers.h b/src/runtime/models/timm_repvgg/plugin_helpers.h new file mode 100644 index 0000000000..bb30ac6982 --- /dev/null +++ b/src/runtime/models/timm_repvgg/plugin_helpers.h @@ -0,0 +1,130 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +// Shared helper functions for pipeline plugins. +// Extracted from pipeline_factory.cpp's anonymous namespace so all +// strategy plugins can reuse TRT module loading, tokenizer creation, +// KV-dim computation, and data-section conversion utilities. + +#include "bundle/bundle_format.h" +#include "bundle/bundle_view.h" +#include "trtmc/runtime/pipeline_plugin.h" +#include "trtmc/runtime/trt_backend.h" +#include "trtmc/runtime/trt_module.h" +#include "trtmc/tokenizer.h" + +#include +#include +#include +#include + +namespace trtmc { + +// A loaded TRT engine, ready for inference. +// The stream is owned internally by the module — callers get it via module->stream(). +struct LoadedModule { + std::unique_ptr module; +}; + +// Load a TRT engine from a serialized plan via the backend. Throws on failure. +LoadedModule load_trt_module_from_plan(IBackend* backend, const std::vector* plan, + const char* label, const ModuleCreateOptions& options = {}); + +// Emit a parseable runtime load/deserialization timing line. +void log_trt_load_timing(const char* label, double load_deserialize_ms, std::size_t plan_bytes); + +// Like load_trt_module_from_plan but returns empty LoadedModule on failure +// instead of throwing (for optional engines). +LoadedModule try_load_trt_module_from_plan(IBackend* backend, const std::vector* plan, + const char* label, + const ModuleCreateOptions& options = {}); + +// Load an optional TRT module, returning nullptr if the plan is absent. +// On deserialization failure, returns nullptr (does not throw). +std::unique_ptr extract_optional_module(IBackend* backend, + const std::vector* plan, + const char* label, + const ModuleCreateOptions& options = {}); + +// Dual-profile TRT module group: one shared backend engine, two module +// contexts (one per optimization profile). Weights live once in GPU memory +// and both modules share the CUDA stream. Use `decode->stream()` to obtain +// the shared stream. +struct DualProfileModules { + std::unique_ptr prefill; // batched Sq profile (null if single-profile) + std::unique_ptr decode; // Sq=1 profile, or the only profile if single-profile +}; + +// Load an engine from a serialized plan via the backend and create two +// execution contexts — one per optimization profile — sharing the engine. +// When the engine has fewer than 2 profiles, `prefill` is left null and +// `decode` holds the single-profile context (legacy bundles). +DualProfileModules load_dual_profile_modules(IBackend* backend, const std::vector* plan, + const char* label, + const ModuleCreateOptions& options = {}); + +// Detect whether the bundle's config requests add_special_tokens for the tokenizer. +bool detect_add_special_tokens(const BundleFile& bundle); + +// Check if the bundle's tokenizer.json describes a BPE model. +bool is_bpe_tokenizer_json(const BundleFile& bundle); + +// Try to create a native C++ BPE tokenizer from the bundle's tokenizer.json. +// Returns nullptr if the section is absent or the model is non-BPE. +// If throw_on_failure is true, throws instead of returning nullptr on BPE parse errors. +std::shared_ptr try_create_native_bpe(const BundleFile& bundle, bool add_special, + bool throw_on_failure); + +// Try to create a native C++ tokenizer from the bundle's tokenizer.json. +// Attempts: BPE -> WordPiece -> Unigram. Returns nullptr if none match. +std::shared_ptr try_create_native_tokenizer(const BundleFile& bundle, + bool add_special_tokens); + +// Create a native tokenizer from bundle. Tries BPE -> WordPiece -> Unigram. +// Returns nullptr if no native tokenizer matches. +std::shared_ptr create_tokenizer_from_bundle(const BundleFile& bundle); + +// Compute the KV cache dimension from model config. +int32_t compute_kv_dim(const BaseConfig& cfg); + +// Convert the BaseConfig precision string ("fp16", "bf16", "fp32") to a DType +// for use as KV cache element type. +DType cache_dtype_from_precision(const std::string& precision); + +// Reinterpret a raw char section as a vector of floats. +std::vector section_to_floats(const std::vector* sec); + +// Reinterpret a raw char section as a vector of int32_t. +std::vector section_to_int32s(const std::vector* sec); + +// Return true if the section pointer is non-null and non-empty. +bool has_section_data(const std::vector* d); + +// BundleFile-based helpers. + +// Mel filterbank loaded from bundle (for Whisper native mel extraction). +struct MelFilterbank { + std::vector data; // [n_freq_bins * n_mel_bins] row-major + int32_t n_freq_bins{0}; + int32_t n_mel_bins{0}; +}; + +// Load mel filterbank from the "mel_filterbank" bundle section. +// Returns empty MelFilterbank if section is not present (old bundles). +MelFilterbank load_mel_filterbank(const BundleFile& bundle); + +// Create a native BPE tokenizer from the CLIP tokenizer sections in the bundle. +// Used for dual-tokenizer models (e.g., FLUX: CLIP + T5). +// Returns nullptr if clip_tokenizer.json section is absent. +std::unique_ptr create_clip_tokenizer_from_bundle(const BundleFile& bundle); + +// Load all TVM-FFI kernels listed in the bundle's kernel_manifest.json. +// Must be called BEFORE deserializing any TRT engine that uses FFI plugins. +// No-op if the bundle has no kernel_manifest.json section (non-FFI bundles). +void load_ffi_kernels_from_bundle(const BundleFile& bundle); + +} // namespace trtmc diff --git a/tests/cpp/models/timm_repvgg/test_timm_repvgg_image_preprocess_seam.cpp b/tests/cpp/models/timm_repvgg/test_timm_repvgg_image_preprocess_seam.cpp new file mode 100644 index 0000000000..673e079446 --- /dev/null +++ b/tests/cpp/models/timm_repvgg/test_timm_repvgg_image_preprocess_seam.cpp @@ -0,0 +1,111 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/models/timm_repvgg/image_preprocess_seam.h" + +#include +#include +#include +#include + +namespace { + +int g_failures = 0; + +void check(bool condition, const char* name) { + if (!condition) { + std::cerr << "FAIL: " << name << '\n'; + ++g_failures; + } +} + +void check_close(float actual, float expected, float tolerance, const char* name) { + if (std::fabs(actual - expected) > tolerance) { + std::cerr << "FAIL: " << name << " actual=" << actual << " expected=" << expected << '\n'; + ++g_failures; + } +} + +void test_timm_repvgg_preprocess_uses_configured_bilinear_resize() { + const std::vector pixels = { + 0.0F, 0.0F, 0.0F, 1.0F, 0.0F, 0.0F, 1.0F, 0.0F, 0.0F, 0.0F, 0.0F, 0.0F, + }; + trtmc::TimmRepvggPreprocessConfig config; + config.input_image_h = 3; + config.input_image_w = 3; + config.crop_pct = 1.0F; + config.interpolation = "bilinear"; + config.image_mean = {0.0F, 0.0F, 0.0F}; + config.image_std = {1.0F, 1.0F, 1.0F}; + + const auto pixel_values = trtmc::preprocess_timm_repvgg_image(pixels.data(), 2, 2, config); + check(pixel_values.size() == 27, "timm RepVGG preprocess size"); + if (pixel_values.size() == 27) { + check_close(pixel_values[4], 0.5F, 1e-6F, + "timm RepVGG bilinear resize blends center pixel"); + } +} + +void test_timm_repvgg_preprocess_applies_bundle_normalization() { + const std::vector pixels(3U * 2U * 2U, 0.75F); + trtmc::TimmRepvggPreprocessConfig config; + config.input_image_h = 2; + config.input_image_w = 2; + config.crop_pct = 1.0F; + config.interpolation = "bilinear"; + config.image_mean = {0.25F, 0.5F, 0.75F}; + config.image_std = {0.5F, 0.25F, 0.125F}; + + const auto pixel_values = trtmc::preprocess_timm_repvgg_image(pixels.data(), 2, 2, config); + check_close(pixel_values[0], 1.0F, 1e-6F, "timm RepVGG red normalization"); + check_close(pixel_values[4], 1.0F, 1e-6F, "timm RepVGG green normalization"); + check_close(pixel_values[8], 0.0F, 1e-6F, "timm RepVGG blue normalization"); +} + +void test_timm_repvgg_resize_matches_torchvision_short_edge_geometry() { + trtmc::TimmRepvggPreprocessConfig config; + config.input_image_h = 224; + config.input_image_w = 224; + config.crop_pct = 0.9F; + + const auto landscape = trtmc::compute_timm_repvgg_resize_shape(320, 426, config); + check(landscape.height == 248, "timm RepVGG landscape short edge uses floor size"); + check(landscape.width == 330, "timm RepVGG landscape aspect ratio uses floor size"); + + const auto portrait = trtmc::compute_timm_repvgg_resize_shape(426, 320, config); + check(portrait.height == 330, "timm RepVGG portrait aspect ratio uses floor size"); + check(portrait.width == 248, "timm RepVGG portrait short edge uses floor size"); +} + +void test_timm_repvgg_preprocess_rejects_invalid_interpolation() { + bool threw = false; + try { + const std::vector pixels(12U, 0.0F); + trtmc::TimmRepvggPreprocessConfig config; + config.input_image_h = 2; + config.input_image_w = 2; + config.crop_pct = 1.0F; + config.interpolation = "nearest"; + (void)trtmc::preprocess_timm_repvgg_image(pixels.data(), 2, 2, config); + } catch (const std::invalid_argument&) { + threw = true; + } + check(threw, "timm RepVGG rejects unsupported interpolation"); +} + +} // namespace + +int main() { + test_timm_repvgg_preprocess_uses_configured_bilinear_resize(); + test_timm_repvgg_preprocess_applies_bundle_normalization(); + test_timm_repvgg_resize_matches_torchvision_short_edge_geometry(); + test_timm_repvgg_preprocess_rejects_invalid_interpolation(); + + if (g_failures != 0) { + std::cerr << g_failures << " timm RepVGG preprocess test(s) failed\n"; + return 1; + } + return 0; +} diff --git a/tests/e2e/models/timm_repvgg/MODEL.toml b/tests/e2e/models/timm_repvgg/MODEL.toml new file mode 100644 index 0000000000..7702163676 --- /dev/null +++ b/tests/e2e/models/timm_repvgg/MODEL.toml @@ -0,0 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +id = "timm_repvgg" +plugin = "timm_repvgg" +test_manifests = [ + "manifests/repvgg-a2-rvgg-in1k.json", +] + +[e2e_defaults.image_classification] +reference_backend = "hf_transformers" +oracle_level = "L1_external_reference" +stages = [ + { name = "full_inference", required = true }, +] +input_fields = [ + { input = "image", manifest = "test_image" }, +] +preflight_asset_fields = ["test_image"] diff --git a/tests/e2e/models/timm_repvgg/data/test_img.jpeg b/tests/e2e/models/timm_repvgg/data/test_img.jpeg new file mode 100644 index 0000000000..095c03620f Binary files /dev/null and b/tests/e2e/models/timm_repvgg/data/test_img.jpeg differ diff --git a/tests/e2e/models/timm_repvgg/e2e_plugins/__init__.py b/tests/e2e/models/timm_repvgg/e2e_plugins/__init__.py new file mode 100644 index 0000000000..1f0515e3ff --- /dev/null +++ b/tests/e2e/models/timm_repvgg/e2e_plugins/__init__.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Model-local E2E plugin package. + +Concrete runner, comparator, and reference implementations are copied into +this package so a model test does not import central concrete E2E strategies. +""" + +from __future__ import annotations + +import os + + +def _case_artifact_dir(artifacts_dir: str, case_name: str) -> str: + if case_name: + d = os.path.join(artifacts_dir, case_name) + else: + d = artifacts_dir + os.makedirs(d, exist_ok=True) + return d + + +def save_full_stderr(stderr: str, artifacts_dir: str, stage_name: str, case_name: str = "") -> tuple: + truncated = stderr[-2000:] if len(stderr) > 2000 else stderr + if not artifacts_dir: + return truncated, None + d = _case_artifact_dir(artifacts_dir, case_name) + path = os.path.join(d, f"{stage_name}_stderr.log") + with open(path, "w", encoding="utf-8") as f: + f.write(stderr) + return truncated, path diff --git a/tests/e2e/models/timm_repvgg/e2e_plugins/benchmark_trt_paths.py b/tests/e2e/models/timm_repvgg/e2e_plugins/benchmark_trt_paths.py new file mode 100644 index 0000000000..45f0da1293 --- /dev/null +++ b/tests/e2e/models/timm_repvgg/e2e_plugins/benchmark_trt_paths.py @@ -0,0 +1,466 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Compare raw TensorRT API and ONNX->trtexec engines for timm RepVGG. + +The raw API path uses TensorRT-Model-Connect's timm_repvgg family plugin. The +ONNX path exports the same timm PyTorch model, builds an engine with trtexec, +and benchmarks both plans with identical input tensors via TensorRT Python. +""" + +from __future__ import annotations + +import argparse +import json +import re +import shutil +import statistics +import subprocess +import sys +import time +from pathlib import Path +from typing import Any + +import numpy as np + + +DEFAULT_MODEL_ID = "timm/repvgg_a2.rvgg_in1k" + + +def _repo_root() -> Path: + for parent in Path(__file__).resolve().parents: + if (parent / "CMakeLists.txt").is_file() and (parent / "python").is_dir(): + return parent + raise RuntimeError("Cannot locate TensorRT-Model-Connect repository root") + + +def _find_trtexec(explicit: str | None) -> str: + candidates = [] + if explicit: + candidates.append(explicit) + which = shutil.which("trtexec") + if which: + candidates.append(which) + candidates.extend([ + "/usr/src/tensorrt/bin/trtexec", + "/opt/tensorrt/bin/trtexec", + "/usr/local/tensorrt/bin/trtexec", + ]) + for candidate in candidates: + if candidate and Path(candidate).is_file(): + return candidate + raise FileNotFoundError( + "trtexec not found; pass --trtexec or add TensorRT's bin directory to PATH" + ) + + +def _create_timm_model(model_id: str): + import timm + + try: + return timm.create_model(f"hf-hub:{model_id}", pretrained=True) + except Exception: + return timm.create_model(f"hf_hub:{model_id}", pretrained=True) + + +def _build_api_engine(model_id: str, plan_path: Path, *, verbose: bool) -> None: + from tensorrt_model_connect.config import ModelConfig + from tensorrt_model_connect.engine_builder import _resolve_model + from tensorrt_model_connect.families import find_plugin + + model_dir = Path(_resolve_model(model_id)) + config = ModelConfig.from_dir(model_dir) + plugin = find_plugin(config.model_type) + if plugin is None or plugin.name != "timm_repvgg": + raise RuntimeError( + f"Expected timm_repvgg plugin for {config.model_type!r}, got {plugin}" + ) + + weights = plugin.load_weights(str(model_dir), config, precision="fp32") + plan = plugin.build_engine( + config, + weights, + max_cache_length=1, + precision="fp32", + verbose=verbose, + ) + plan_path.write_bytes(plan) + + +def _export_onnx(model_id: str, onnx_path: Path) -> None: + import torch + + model = _create_timm_model(model_id) + model.eval() + dummy = torch.randn(1, 3, 224, 224, dtype=torch.float32) + torch.onnx.export( + model, + dummy, + str(onnx_path), + input_names=["pixel_values"], + output_names=["logits"], + opset_version=17, + do_constant_folding=True, + dynamo=False, + ) + + +def _build_trtexec_engine( + trtexec: str, + onnx_path: Path, + plan_path: Path, + log_path: Path, +) -> None: + cmd = [ + trtexec, + f"--onnx={onnx_path}", + f"--saveEngine={plan_path}", + "--skipInference", + ] + result = subprocess.run( + cmd, capture_output=True, text=True, timeout=1800) + log_path.write_text( + "COMMAND: " + " ".join(cmd) + "\n\nSTDOUT:\n" + result.stdout + + "\n\nSTDERR:\n" + result.stderr, + encoding="utf-8", + ) + if result.returncode != 0: + raise RuntimeError( + f"trtexec failed with rc={result.returncode}; see {log_path}") + + +_TRTEXEC_TIMING_RE = re.compile( + r"GPU Compute Time: min = (?P[0-9.]+) ms, max = (?P[0-9.]+) ms, " + r"mean = (?P[0-9.]+) ms, median = (?P[0-9.]+) ms" +) + + +def _benchmark_plan_with_trtexec( + trtexec: str, + plan_path: Path, + log_path: Path, + *, + warmup: int, + iterations: int, +) -> dict[str, Any]: + cmd = [ + trtexec, + f"--loadEngine={plan_path}", + f"--warmUp={max(200, warmup)}", + "--duration=0", + f"--iterations={iterations}", + "--noDataTransfers", + ] + result = subprocess.run( + cmd, capture_output=True, text=True, timeout=1800) + log_path.write_text( + "COMMAND: " + " ".join(cmd) + "\n\nSTDOUT:\n" + result.stdout + + "\n\nSTDERR:\n" + result.stderr, + encoding="utf-8", + ) + if result.returncode != 0: + raise RuntimeError( + f"trtexec benchmark failed with rc={result.returncode}; see {log_path}") + matches = list(_TRTEXEC_TIMING_RE.finditer(result.stdout + "\n" + result.stderr)) + if not matches: + raise RuntimeError(f"Could not parse trtexec GPU timing from {log_path}") + timing = matches[-1].groupdict() + return { + "mean_ms": float(timing["mean"]), + "median_ms": float(timing["median"]), + "stdev_ms": None, + "min_ms": float(timing["min"]), + "max_ms": float(timing["max"]), + "iterations": iterations, + "warmup": max(200, warmup), + "log": str(log_path), + } + + +def _input_from_image(image_path: Path) -> np.ndarray: + from PIL import Image + + target = 224 + crop_pct = 0.9 + resize_short = int(target / crop_pct + 0.5) + image = Image.open(image_path).convert("RGB") + width, height = image.size + if height <= width: + resized_h = resize_short + resized_w = max(1, int(width * resize_short / height + 0.5)) + else: + resized_w = resize_short + resized_h = max(1, int(height * resize_short / width + 0.5)) + image = image.resize((resized_w, resized_h), Image.Resampling.NEAREST) + left = max(0, (resized_w - target) // 2) + top = max(0, (resized_h - target) // 2) + image = image.crop((left, top, left + target, top + target)) + arr = np.asarray(image, dtype=np.float32) / 255.0 + arr = (arr - 0.5) / 0.5 + return np.transpose(arr, (2, 0, 1))[None, ...].copy() + + +def _make_input(seed: int, image: str | None) -> np.ndarray: + if image: + return _input_from_image(Path(image)) + rng = np.random.default_rng(seed) + return rng.normal(0.0, 1.0, size=(1, 3, 224, 224)).astype(np.float32) + + +def _torch_dtype_for_trt(trt_dtype): + import tensorrt as trt + import torch + + mapping = { + trt.float32: torch.float32, + trt.float16: torch.float16, + trt.int32: torch.int32, + trt.bool: torch.bool, + } + if hasattr(trt, "bfloat16"): + mapping[trt.bfloat16] = torch.bfloat16 + if hasattr(trt, "int64"): + mapping[trt.int64] = torch.int64 + return mapping[trt_dtype] + + +def _load_engine(plan_path: Path): + import tensorrt as trt + + logger = trt.Logger(trt.Logger.WARNING) + runtime = trt.Runtime(logger) + engine = runtime.deserialize_cuda_engine(plan_path.read_bytes()) + if engine is None: + raise RuntimeError(f"Failed to deserialize {plan_path}") + return engine + + +def _prepare_context(plan_path: Path, input_np: np.ndarray): + import tensorrt as trt + import torch + + engine = _load_engine(plan_path) + context = engine.create_execution_context() + tensors: dict[str, Any] = {} + outputs: dict[str, Any] = {} + + for idx in range(engine.num_io_tensors): + name = engine.get_tensor_name(idx) + mode = engine.get_tensor_mode(name) + if mode == trt.TensorIOMode.INPUT: + if hasattr(context, "set_input_shape"): + context.set_input_shape(name, tuple(input_np.shape)) + tensor = torch.as_tensor(input_np, device="cuda").contiguous() + else: + shape = tuple(int(dim) for dim in context.get_tensor_shape(name)) + dtype = _torch_dtype_for_trt(engine.get_tensor_dtype(name)) + tensor = torch.empty(shape, dtype=dtype, device="cuda") + outputs[name] = tensor + tensors[name] = tensor + context.set_tensor_address(name, int(tensor.data_ptr())) + + return engine, context, tensors, outputs + + +def _run_once(context, outputs: dict[str, Any], stream) -> np.ndarray: + import torch + + with torch.cuda.stream(stream): + context.execute_async_v3(stream_handle=stream.cuda_stream) + torch.cuda.synchronize() + if "logits" in outputs: + out = outputs["logits"] + else: + out = next(iter(outputs.values())) + return out.detach().float().cpu().numpy().reshape(-1) + + +def _benchmark_plan( + plan_path: Path, + input_np: np.ndarray, + *, + warmup: int, + iterations: int, +) -> dict[str, Any]: + import torch + + _engine, context, _tensors, outputs = _prepare_context(plan_path, input_np) + stream = torch.cuda.Stream() + for _ in range(warmup): + with torch.cuda.stream(stream): + context.execute_async_v3(stream_handle=stream.cuda_stream) + torch.cuda.synchronize() + + times_ms: list[float] = [] + for _ in range(iterations): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + with torch.cuda.stream(stream): + start.record(stream) + context.execute_async_v3(stream_handle=stream.cuda_stream) + end.record(stream) + torch.cuda.synchronize() + times_ms.append(float(start.elapsed_time(end))) + + logits = _run_once(context, outputs, stream) + return { + "mean_ms": float(statistics.fmean(times_ms)), + "median_ms": float(statistics.median(times_ms)), + "stdev_ms": float(statistics.pstdev(times_ms)), + "min_ms": float(min(times_ms)), + "max_ms": float(max(times_ms)), + "iterations": iterations, + "warmup": warmup, + "logits": logits, + } + + +def _summarize( + api: dict[str, Any], + onnx: dict[str, Any], + *, + max_ratio: float, + max_abs_diff: float, +) -> dict[str, Any]: + api_logits = api.pop("logits") + onnx_logits = onnx.pop("logits") + abs_diff = np.abs(api_logits - onnx_logits) + top_api = int(np.argmax(api_logits)) + top_onnx = int(np.argmax(onnx_logits)) + api_over_onnx = api["mean_ms"] / onnx["mean_ms"] + symmetric_ratio = ( + max(api["mean_ms"], onnx["mean_ms"]) / min(api["mean_ms"], onnx["mean_ms"]) + ) + return { + "api_engine": api, + "onnx_trtexec_engine": onnx, + "api_over_onnx_ratio": float(api_over_onnx), + "perf_ratio_max_over_min": float(symmetric_ratio), + "perf_within_threshold": bool(api_over_onnx <= max_ratio), + "max_allowed_perf_ratio": max_ratio, + "output": { + "max_abs_diff": float(abs_diff.max()), + "mean_abs_diff": float(abs_diff.mean()), + "top1_api": top_api, + "top1_onnx_trtexec": top_onnx, + "top1_match": top_api == top_onnx, + "max_allowed_abs_diff": max_abs_diff, + "within_threshold": bool(abs_diff.max() <= max_abs_diff), + }, + } + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-id", default=DEFAULT_MODEL_ID) + parser.add_argument( + "--out-dir", + default=str(_repo_root() / "artifacts" / "timm_repvgg_trt_path_comparison"), + ) + parser.add_argument("--trtexec", default=None) + parser.add_argument("--image", default=None) + parser.add_argument("--iterations", type=int, default=100) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--seed", type=int, default=1234) + parser.add_argument("--max-ratio", type=float, default=1.05) + parser.add_argument("--max-abs-diff", type=float, default=0.05) + parser.add_argument( + "--perf-runner", + choices=("trtexec", "python"), + default="trtexec", + help=( + "Use direct trtexec GPU Compute Time for the performance gate " + "or Python CUDA events. Python is still used for correctness." + ), + ) + parser.add_argument("--rebuild", action="store_true") + parser.add_argument("--verbose", action="store_true") + return parser.parse_args() + + +def main() -> int: + args = parse_args() + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + api_plan = out_dir / "api.plan" + onnx_path = out_dir / "model.onnx" + onnx_plan = out_dir / "onnx_trtexec.plan" + trtexec_log = out_dir / "trtexec.log" + result_path = out_dir / "result.json" + + trtexec = _find_trtexec(args.trtexec) + + t0 = time.monotonic() + if args.rebuild or not api_plan.is_file(): + _build_api_engine(args.model_id, api_plan, verbose=args.verbose) + if args.rebuild or not onnx_path.is_file(): + _export_onnx(args.model_id, onnx_path) + if args.rebuild or not onnx_plan.is_file(): + _build_trtexec_engine(trtexec, onnx_path, onnx_plan, trtexec_log) + + input_np = _make_input(args.seed, args.image) + api_python = _benchmark_plan( + api_plan, input_np, warmup=args.warmup, iterations=args.iterations) + onnx_python = _benchmark_plan( + onnx_plan, input_np, warmup=args.warmup, iterations=args.iterations) + if args.perf_runner == "trtexec": + api_perf = _benchmark_plan_with_trtexec( + trtexec, + api_plan, + out_dir / "trtexec_api_bench.log", + warmup=args.warmup, + iterations=args.iterations, + ) + onnx_perf = _benchmark_plan_with_trtexec( + trtexec, + onnx_plan, + out_dir / "trtexec_onnx_bench.log", + warmup=args.warmup, + iterations=args.iterations, + ) + api = {**api_perf, "logits": api_python["logits"]} + onnx = {**onnx_perf, "logits": onnx_python["logits"]} + else: + api = api_python.copy() + onnx = onnx_python.copy() + summary = _summarize( + api, onnx, max_ratio=args.max_ratio, max_abs_diff=args.max_abs_diff) + summary["perf_runner"] = args.perf_runner + summary["python_api_engine"] = { + key: value for key, value in api_python.items() if key != "logits"} + summary["python_onnx_trtexec_engine"] = { + key: value for key, value in onnx_python.items() if key != "logits"} + summary.update({ + "model_id": args.model_id, + "api_plan": str(api_plan), + "onnx_path": str(onnx_path), + "onnx_trtexec_plan": str(onnx_plan), + "trtexec": trtexec, + "trtexec_log": str(trtexec_log), + "input": {"image": args.image, "seed": args.seed}, + "total_elapsed_s": time.monotonic() - t0, + }) + result_path.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8") + + print( + "api_mean_ms={api:.6f} onnx_trtexec_mean_ms={onnx:.6f} " + "api_over_onnx={ratio:.6f} top1_match={top1} max_abs_diff={diff:.6g} " + "result={result}".format( + api=summary["api_engine"]["mean_ms"], + onnx=summary["onnx_trtexec_engine"]["mean_ms"], + ratio=summary["api_over_onnx_ratio"], + top1=summary["output"]["top1_match"], + diff=summary["output"]["max_abs_diff"], + result=result_path, + ) + ) + if not summary["perf_within_threshold"]: + return 2 + if not summary["output"]["top1_match"] or not summary["output"]["within_threshold"]: + return 3 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/e2e/models/timm_repvgg/e2e_plugins/comparator.py b/tests/e2e/models/timm_repvgg/e2e_plugins/comparator.py new file mode 100644 index 0000000000..dc4edb61a3 --- /dev/null +++ b/tests/e2e/models/timm_repvgg/e2e_plugins/comparator.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""timm_repvgg model-owned E2E comparator plugins.""" + +from __future__ import annotations + +from .comparators.image_classification import ImageClassificationComparator + + +class TimmRepvggImageClassificationComparator(ImageClassificationComparator): + """timm_repvgg local comparator for image_classification.""" + +comparator = TimmRepvggImageClassificationComparator() diff --git a/tests/e2e/models/timm_repvgg/e2e_plugins/comparators/__init__.py b/tests/e2e/models/timm_repvgg/e2e_plugins/comparators/__init__.py new file mode 100644 index 0000000000..9c9857522e --- /dev/null +++ b/tests/e2e/models/timm_repvgg/e2e_plugins/comparators/__init__.py @@ -0,0 +1,9 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Comparators — TRT vs reference output comparison for each task strategy. + +Each module in this package should expose a module-level ``plugin`` attribute +that is an instance implementing the Comparator protocol. The registry +auto-discovers these plugins on first access. +""" diff --git a/tests/e2e/models/timm_repvgg/e2e_plugins/comparators/_helpers.py b/tests/e2e/models/timm_repvgg/e2e_plugins/comparators/_helpers.py new file mode 100644 index 0000000000..20d4d897b6 --- /dev/null +++ b/tests/e2e/models/timm_repvgg/e2e_plugins/comparators/_helpers.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared helper functions for comparator modules. + +These utilities were previously duplicated across multiple comparator files. +The canonical implementations come from text.py. +""" + +from __future__ import annotations + +import numpy as np + + +def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float: + """Cosine similarity between two 1-D vectors. Returns 0.0 on degenerate input.""" + norm_a = np.linalg.norm(a) + norm_b = np.linalg.norm(b) + if norm_a < 1e-12 or norm_b < 1e-12: + return 0.0 + return float(np.dot(a, b) / (norm_a * norm_b)) + + +def levenshtein_distance(s1: str, s2: str) -> int: + """Standard Levenshtein edit distance via dynamic programming.""" + if len(s1) < len(s2): + return levenshtein_distance(s2, s1) + if len(s2) == 0: + return len(s1) + + prev_row = list(range(len(s2) + 1)) + for i, c1 in enumerate(s1): + curr_row = [i + 1] + for j, c2 in enumerate(s2): + insertions = prev_row[j + 1] + 1 + deletions = curr_row[j] + 1 + substitutions = prev_row[j] + (0 if c1 == c2 else 1) + curr_row.append(min(insertions, deletions, substitutions)) + prev_row = curr_row + return prev_row[-1] + + +def normalized_edit_distance(s1: str, s2: str) -> float: + """Levenshtein distance normalized by max string length. 0.0 = identical.""" + max_len = max(len(s1), len(s2)) + if max_len == 0: + return 0.0 + return levenshtein_distance(s1, s2) / max_len diff --git a/tests/e2e/models/timm_repvgg/e2e_plugins/comparators/image_classification.py b/tests/e2e/models/timm_repvgg/e2e_plugins/comparators/image_classification.py new file mode 100644 index 0000000000..51ebd87e73 --- /dev/null +++ b/tests/e2e/models/timm_repvgg/e2e_plugins/comparators/image_classification.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Image-classification comparator.""" + +from __future__ import annotations + +from ..contracts import ( + CompareResult, + MetricResult, + StageOutput, + StageSpec, + StageStatus, + ThresholdProfile, +) + + +class ImageClassificationComparator: + @property + def task_strategy(self) -> str: + return "image_classification" + + def compare( + self, + trt: StageOutput, + ref: StageOutput, + threshold: ThresholdProfile, + stage: StageSpec, + ) -> CompareResult: + metrics: dict[str, MetricResult] = {} + + trt_top = trt.data.get("top_class") + ref_top = ref.data.get("top_class") + if trt_top is None: + return CompareResult( + stage_name=stage.name, + status=StageStatus.ERROR.value, + metrics=metrics, + message="TRT classification output missing top_class", + ) + if ref_top is None: + return CompareResult( + stage_name=stage.name, + status=StageStatus.ERROR.value, + metrics=metrics, + message="Reference classification output missing top_class", + ) + + top1_match = int(trt_top) == int(ref_top) + metrics["top1_match"] = MetricResult( + value=1.0 if top1_match else 0.0, + threshold=1.0, + operator="==", + passed=top1_match, + ) + + if "top_score" in trt.data and "top_score" in ref.data: + diff = abs(float(trt.data["top_score"]) - float(ref.data["top_score"])) + score_atol = threshold.metrics.get("top_score_atol") + metrics["top_score_abs_diff"] = MetricResult( + value=diff, + threshold=score_atol, + operator="<=" if score_atol is not None else "informational", + passed=True if score_atol is None else diff <= score_atol, + ) + + passed = all(metric.passed for metric in metrics.values()) + return CompareResult( + stage_name=stage.name, + status=StageStatus.PASSED.value if passed else StageStatus.FAILED.value, + metrics=metrics, + composite_rule="top-1 class must match", + message=( + f"Image classification: TRT top={int(trt_top)}, " + f"reference top={int(ref_top)}" + ), + ) + + +plugin = ImageClassificationComparator() diff --git a/tests/e2e/models/timm_repvgg/e2e_plugins/contract.py b/tests/e2e/models/timm_repvgg/e2e_plugins/contract.py new file mode 100644 index 0000000000..f7bcfb23d0 --- /dev/null +++ b/tests/e2e/models/timm_repvgg/e2e_plugins/contract.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""TIMM ViT-owned image classification contract plugin.""" + +from __future__ import annotations + +from tests.e2e_harness.contracts import CompareResult, MetricResult + + +def _make_pass(stage_name: str, metrics, rule: str) -> CompareResult: + return CompareResult( + stage_name=stage_name, + status="passed", + metrics=metrics, + composite_rule=rule, + message="TIMM ViT image classification contract verified", + ) + + +def _make_fail(stage_name: str, metrics, rule: str, message: str) -> CompareResult: + return CompareResult( + stage_name=stage_name, + status="failed", + metrics=metrics, + composite_rule=rule, + message=message, + ) + + +class TimmRepvggImageClassificationPlugin: + reference_families = ["image_classification"] + user_contract = "image_classification" + + def configure_reference(self, case): + del case + return {} + + def verify(self, trt_output, ref_output, case, threshold): + del case + stage = trt_output.stage_name or "full_inference" + metrics: dict[str, MetricResult] = {} + trt_top = trt_output.data.get("top_class") + ref_top = ref_output.data.get("top_class") + if trt_top is None: + return _make_fail(stage, metrics, "top-1 class must match", "TRT output missing top_class") + if ref_top is None: + return _make_fail( + stage, + metrics, + "top-1 class must match", + "Reference output missing top_class", + ) + + top1_match = int(trt_top) == int(ref_top) + metrics["top1_match"] = MetricResult( + value=1.0 if top1_match else 0.0, + threshold=1.0, + operator="==", + passed=top1_match, + ) + + if "top_score" in trt_output.data and "top_score" in ref_output.data: + diff = abs(float(trt_output.data["top_score"]) - float(ref_output.data["top_score"])) + score_atol = threshold.metrics.get("top_score_atol") + metrics["top_score_abs_diff"] = MetricResult( + value=diff, + threshold=score_atol, + operator="<=" if score_atol is not None else "informational", + passed=True if score_atol is None else diff <= score_atol, + ) + + passed = all(metric.passed for metric in metrics.values()) + rule = "top-1 class must match" + if passed: + return _make_pass(stage, metrics, rule) + return _make_fail( + stage, + metrics, + rule, + f"TIMM ViT classification mismatch: TRT top={trt_top}, reference top={ref_top}", + ) + + +plugin = TimmRepvggImageClassificationPlugin() diff --git a/tests/e2e/models/timm_repvgg/e2e_plugins/contracts.py b/tests/e2e/models/timm_repvgg/e2e_plugins/contracts.py new file mode 100644 index 0000000000..d6f9281d2a --- /dev/null +++ b/tests/e2e/models/timm_repvgg/e2e_plugins/contracts.py @@ -0,0 +1,10 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Model-local E2E contract aliases. + +Contracts remain the stable harness API; concrete runners/references/comparators +are owned by the model package. +""" + +from tests.e2e_harness.contracts import * # noqa: F401,F403 diff --git a/tests/e2e/models/timm_repvgg/e2e_plugins/reference.py b/tests/e2e/models/timm_repvgg/e2e_plugins/reference.py new file mode 100644 index 0000000000..5bf9599153 --- /dev/null +++ b/tests/e2e/models/timm_repvgg/e2e_plugins/reference.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""timm_repvgg model-owned E2E reference plugins.""" + +from __future__ import annotations + +from .references.hf_transformers import HfTransformersReference + + +class TimmRepvggHfTransformersReference(HfTransformersReference): + """timm_repvgg local reference for hf_transformers.""" + +reference = TimmRepvggHfTransformersReference() diff --git a/tests/e2e/models/timm_repvgg/e2e_plugins/references/__init__.py b/tests/e2e/models/timm_repvgg/e2e_plugins/references/__init__.py new file mode 100644 index 0000000000..6487ae8000 --- /dev/null +++ b/tests/e2e/models/timm_repvgg/e2e_plugins/references/__init__.py @@ -0,0 +1,9 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Reference backends — reference inference for comparison. + +Each module in this package should expose a module-level ``plugin`` attribute +that is an instance implementing the ReferenceBackendRunner protocol. The +registry auto-discovers these plugins on first access. +""" diff --git a/tests/e2e/models/timm_repvgg/e2e_plugins/references/custom_python.py b/tests/e2e/models/timm_repvgg/e2e_plugins/references/custom_python.py new file mode 100644 index 0000000000..e0a861eef7 --- /dev/null +++ b/tests/e2e/models/timm_repvgg/e2e_plugins/references/custom_python.py @@ -0,0 +1,145 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Custom Python reference backend — execute a user-provided Python script. + +The script path is specified in case metadata as ``custom_python_script``. +Flexible for non-standard models that don't fit HF Transformers/Diffusers. +""" + +from __future__ import annotations + +import json +import logging +import os +import subprocess +import time + +from .. import save_full_stderr +from ..contracts import E2ECase, RunContext, StageOutput, StageSpec + +logger = logging.getLogger(__name__) + + +class CustomPythonReference: + """Execute a custom Python script as reference backend.""" + + @property + def backend_name(self) -> str: + return "custom_python" + + def run_stage( + self, case: E2ECase, stage: StageSpec, ctx: RunContext + ) -> StageOutput: + script_path = case.metadata.get("custom_python_script") + if not script_path: + raise ValueError( + f"Case {case.name} uses custom_python reference but " + f"metadata.custom_python_script is not set" + ) + + # Resolve script path relative to project root + if not os.path.isabs(script_path): + project_root = os.path.dirname( + os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + ) + script_path = os.path.join(project_root, script_path) + + python = ctx.reference_python_path() or "python3" + + cmd = [ + python, script_path, + "--model", case.hf_id, + "--stage", stage.name, + ] + + # Pass prompt if available + prompt = case.inputs.get("prompt", "") + if prompt: + cmd.extend(["--prompt", prompt]) + + # Pass image if available + image = case.inputs.get("image") + if image: + cmd.extend(["--image", image]) + + # Pass audio if available + audio = case.inputs.get("audio") + if audio: + cmd.extend(["--audio", audio]) + + max_new_tokens = case.inputs.get("max_new_tokens", 30) + cmd.extend(["--max-new-tokens", str(max_new_tokens)]) + + # Pass trust_remote_code if needed + if case.metadata.get("trust_remote_code"): + cmd.append("--trust-remote-code") + + # Pass any extra script args from metadata + extra_args = case.metadata.get("custom_python_args", []) + cmd.extend(extra_args) + + env = dict(os.environ) + if ctx.ld_library_path: + env["LD_LIBRARY_PATH"] = ctx.ld_library_path + + logger.info("Running custom Python reference: %s", " ".join(cmd)) + t0 = time.monotonic() + result = subprocess.run( + cmd, capture_output=True, text=True, env=env, timeout=1800, + ) + elapsed = time.monotonic() - t0 + + if result.returncode != 0: + truncated, log_path = save_full_stderr( + result.stderr, ctx.artifacts_dir or "", + "custom_python", case.name) + msg = (f"Custom Python reference failed (rc={result.returncode}): " + f"{truncated}") + if log_path: + msg += f" (full stderr: {log_path})" + raise RuntimeError(msg) + + # Parse output: expect JSON on stdout + data = _parse_output(result.stdout.strip()) + + return StageOutput( + stage_name=stage.name, + data=data, + text=data.get("text"), + timing_s=elapsed, + metadata={ + "command": cmd, + "returncode": result.returncode, + "script_path": script_path, + }, + ) + + +def _parse_output(stdout: str) -> dict: + """Parse reference output from custom Python script stdout.""" + # Try JSON (preferred format) + try: + data = json.loads(stdout) + if isinstance(data, dict): + return data + except (json.JSONDecodeError, ValueError): + pass + + # Try parsing last line as JSON (script may print progress before JSON) + for line in reversed(stdout.splitlines()): + line = line.strip() + if not line: + continue + try: + data = json.loads(line) + if isinstance(data, dict): + return data + except (json.JSONDecodeError, ValueError): + continue + + # Fall back to raw text output + return {"text": stdout, "raw_output": stdout} + + +plugin = CustomPythonReference() diff --git a/tests/e2e/models/timm_repvgg/e2e_plugins/references/golden_snapshot.py b/tests/e2e/models/timm_repvgg/e2e_plugins/references/golden_snapshot.py new file mode 100644 index 0000000000..bf2ad85358 --- /dev/null +++ b/tests/e2e/models/timm_repvgg/e2e_plugins/references/golden_snapshot.py @@ -0,0 +1,129 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Golden snapshot reference backend — load pre-computed reference outputs. + +Loads golden outputs from a trusted prior run. The snapshot path is +specified in case metadata as ``golden_snapshot_path`` (directory or file). +""" + +from __future__ import annotations + +import json +import logging +import os +from typing import Any + +from ..contracts import E2ECase, RunContext, StageOutput, StageSpec + +logger = logging.getLogger(__name__) + + +class GoldenSnapshotReference: + """Load pre-computed golden outputs as reference.""" + + @property + def backend_name(self) -> str: + return "golden_snapshot" + + def run_stage( + self, case: E2ECase, stage: StageSpec, ctx: RunContext + ) -> StageOutput: + snapshot_path = case.metadata.get("golden_snapshot_path") + if not snapshot_path: + raise ValueError( + f"Case {case.name} uses golden_snapshot reference but " + f"metadata.golden_snapshot_path is not set" + ) + + # Resolve relative paths against engine dir or project root + if not os.path.isabs(snapshot_path): + if ctx.engine_dir and os.path.exists( + os.path.join(ctx.engine_dir, snapshot_path) + ): + snapshot_path = os.path.join(ctx.engine_dir, snapshot_path) + else: + project_root = os.path.dirname( + os.path.dirname( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + ) + ) + snapshot_path = os.path.join(project_root, snapshot_path) + + data = _load_snapshot(snapshot_path, stage.name) + + return StageOutput( + stage_name=stage.name, + data=data, + text=data.get("text"), + timing_s=0.0, + metadata={ + "source": "golden_snapshot", + "snapshot_path": snapshot_path, + }, + ) + + +def _load_snapshot(snapshot_path: str, stage_name: str) -> dict[str, Any]: + """Load golden snapshot data. + + Supports: + - Directory: looks for .json or .npy + - JSON file: loads directly + - NPY file: loads numpy array as 'output_field' + """ + if os.path.isdir(snapshot_path): + # Look for stage-specific file + json_path = os.path.join(snapshot_path, f"{stage_name}.json") + npy_path = os.path.join(snapshot_path, f"{stage_name}.npy") + + if os.path.isfile(json_path): + return _load_json(json_path) + elif os.path.isfile(npy_path): + return _load_npy(npy_path) + + # Try generic output files + for name in ("output.json", "golden.json", "reference.json"): + path = os.path.join(snapshot_path, name) + if os.path.isfile(path): + return _load_json(path) + + raise FileNotFoundError( + f"No golden snapshot found for stage {stage_name} in {snapshot_path}" + ) + + elif snapshot_path.endswith(".json"): + return _load_json(snapshot_path) + + elif snapshot_path.endswith(".npy") or snapshot_path.endswith(".npz"): + return _load_npy(snapshot_path) + + else: + raise ValueError(f"Unsupported golden snapshot format: {snapshot_path}") + + +def _load_json(path: str) -> dict[str, Any]: + """Load JSON golden snapshot.""" + with open(path, encoding="utf-8") as f: + data = json.load(f) + if isinstance(data, dict): + return data + return {"data": data} + + +def _load_npy(path: str) -> dict[str, Any]: + """Load numpy golden snapshot.""" + try: + import numpy as np + except ImportError: + raise ImportError("numpy is required to load .npy golden snapshots") + + if path.endswith(".npz"): + loaded = np.load(path) + return {key: loaded[key] for key in loaded.files} + else: + arr = np.load(path, allow_pickle=False) + return {"output_field": arr} + + +plugin = GoldenSnapshotReference() diff --git a/tests/e2e/models/timm_repvgg/e2e_plugins/references/hf_transformers.py b/tests/e2e/models/timm_repvgg/e2e_plugins/references/hf_transformers.py new file mode 100644 index 0000000000..bb818c6d2e --- /dev/null +++ b/tests/e2e/models/timm_repvgg/e2e_plugins/references/hf_transformers.py @@ -0,0 +1,1053 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""HuggingFace Transformers reference backend — gold-standard L1 oracle. + +Runs HF model inference in a subprocess for GPU memory isolation and returns +per-step logits + generated text for comparison against TRT outputs. +""" + +from __future__ import annotations + +import json +import logging +import os +import subprocess +import sys +import tempfile +import textwrap +import time +from collections.abc import Callable, Iterable, Mapping, Sequence +from pathlib import Path +from typing import Any + +from .. import save_full_stderr, _case_artifact_dir +from ..contracts import E2ECase, RunContext, StageOutput, StageSpec + +logger = logging.getLogger(__name__) + +PROJECT_DIR = Path(__file__).resolve().parents[6] +E2E_DIR = PROJECT_DIR / "tests" / "e2e" + + +_PRECISION_TO_TORCH_DTYPE = { + "fp16": "torch.float16", + "fp32": "torch.float32", + "bf16": "torch.bfloat16", +} + + +def _torch_dtype_for_case(case: E2ECase) -> str: + """Return the explicit reference dtype, falling back to DUT precision. + + FP16 acceptance manifests set reference_precision=fp32 so changing the + engine precision does not also change the oracle. + """ + precision = case.metadata.get( + "reference_precision", case.metadata.get("precision", "fp32")) + return _PRECISION_TO_TORCH_DTYPE.get(precision, "torch.float32") + + +def _vl_prompt_has_image_placeholder(text: str) -> bool: + """Return true when a rendered VL prompt still carries an image placeholder.""" + return any(marker in text for marker in ( + "<|image_pad|>", + "<|vision_start|>", + "", + "", + )) + + +def _normalize_vl_prompt_guard(text: str) -> str: + """Normalize decoded VL text for prompt-only reference detection.""" + normalized = " ".join(str(text or "").split()).strip().lower() + for marker in ( + "", + "", + "<|image_pad|>", + "<|vision_start|>", + "<|vision_end|>", + ): + normalized = normalized.replace(marker, " ") + return " ".join(normalized.split()).strip() + + +def _is_prompt_only_vl_text(text: str, prompt_texts: tuple[str, ...]) -> bool: + """Return true when decoded VL text contains only the input prompt/template.""" + normalized_text = _normalize_vl_prompt_guard(text) + if not normalized_text: + return True + + for prompt_text in prompt_texts: + normalized_prompt = _normalize_vl_prompt_guard(prompt_text) + if not normalized_prompt: + continue + if normalized_text == normalized_prompt: + return True + if normalized_text.startswith(normalized_prompt): + tail = normalized_text[len(normalized_prompt):].strip(" :") + if tail in {"", "assistant", "answer"}: + return True + if normalized_text.endswith(normalized_prompt): + return True + return False + + +def _decode_vl_generated_text( + processor, + generated_ids, + input_len: int, + prompt_texts: tuple[str, ...] = (), +) -> str: + """Decode VL generation whether generate() returns full ids or generated ids only.""" + token_count = len(generated_ids) + + def _decode_token_ids(token_ids) -> str: + return processor.decode(token_ids, skip_special_tokens=True).strip() + + if input_len > 0 and token_count > input_len: + text = _decode_token_ids(generated_ids[input_len:]) + if text and not _is_prompt_only_vl_text(text, prompt_texts): + return text + + text = _decode_token_ids(generated_ids) + if text and not _is_prompt_only_vl_text(text, prompt_texts): + return text + return "" + + +def _resolve_cached_model_ref(hf_id: str) -> str: + """Prefer a locally cached HF snapshot to avoid Hub API rate limits.""" + if not hf_id: + return hf_id + p = Path(hf_id) + if p.exists(): + return hf_id + + try: + from huggingface_hub import snapshot_download + + return snapshot_download(hf_id, local_files_only=True) + except Exception: + return hf_id + + +ReferenceOutputReader = Callable[[], dict[str, Any]] + + +def _coerce_stream_text(stream: object) -> str: + if stream is None: + return "" + if isinstance(stream, bytes): + return stream.decode(errors="replace") + return str(stream) + + +def _read_text_artifact(path: str, *, encoding: str = "utf-8") -> str: + artifact_path = Path(path) + if not artifact_path.is_file(): + return "" + return artifact_path.read_text(encoding=encoding) + + +def _json_output_reader(path: str, *, encoding: str = "utf-8") -> ReferenceOutputReader: + def _reader() -> dict[str, Any]: + artifact_path = Path(path) + if not artifact_path.is_file(): + return {} + return json.loads(artifact_path.read_text(encoding=encoding)) + + return _reader + + +def _json_text_reader( + path: str, key: str = "text", *, encoding: str = "utf-8" +) -> Callable[[], str]: + def _reader() -> str: + data = _json_output_reader(path, encoding=encoding)() + value = data.get(key, "") + return "" if value is None else str(value) + + return _reader + + +def _npy_output_reader( + path: str, + data_key: str, + *, + path_key: str = "", + allow_pickle: bool = False, +) -> ReferenceOutputReader: + def _reader() -> dict[str, Any]: + artifact_path = Path(path) + if not artifact_path.is_file(): + return {} + import numpy as np + + data: dict[str, Any] = {} + if path_key: + data[path_key] = path + data[data_key] = np.load(artifact_path, allow_pickle=allow_pickle) + return data + + return _reader + + +def _existing_path_reader(path: str, data_key: str) -> ReferenceOutputReader: + def _reader() -> dict[str, Any]: + return {data_key: path} if Path(path).is_file() else {} + + return _reader + + +def _reference_env(ctx: RunContext) -> dict[str, str]: + env = dict(os.environ) + if ctx.ld_library_path: + env["LD_LIBRARY_PATH"] = ctx.ld_library_path + return env + + +def run_reference_subprocess( + *, + command: Sequence[str], + timeout_s: float, + label: str, + artifact_dir: str, + case_name: str, + stage_name: str, + env: Mapping[str, str] | None = None, + output_readers: Iterable[ReferenceOutputReader] = (), + text_reader: Callable[[], str] | None = None, + logits_reader: Callable[[], Any] | None = None, + metadata: Mapping[str, Any] | None = None, + include_stdio_metadata: bool = False, + failure_label: str | None = None, +) -> StageOutput: + """Run a reference subprocess and build the matching StageOutput.""" + failure_prefix = failure_label or label.replace("_", " ") + cmd = list(command) + t0 = time.monotonic() + try: + result = subprocess.run( + cmd, + capture_output=True, + text=True, + timeout=timeout_s, + env=dict(env) if env is not None else None, + ) + except subprocess.TimeoutExpired as exc: + elapsed = time.monotonic() - t0 + stderr = _coerce_stream_text(exc.stderr) + truncated, log_path = save_full_stderr( + stderr, artifact_dir, label, case_name + ) + msg = f"{failure_prefix} timed out for {case_name} after {elapsed:.0f}s" + if truncated: + msg += f":\n{truncated}" + if log_path: + msg += f" (full stderr: {log_path})" + raise RuntimeError(msg) from exc + except Exception as exc: + raise RuntimeError(f"{failure_prefix} failed for {case_name}: {exc}") from exc + elapsed = time.monotonic() - t0 + + if result.returncode != 0: + truncated, log_path = save_full_stderr( + result.stderr or "", artifact_dir, label, case_name + ) + msg = ( + f"{failure_prefix} failed for {case_name} " + f"(rc={result.returncode}):\n{truncated}" + ) + if log_path: + msg += f" (full stderr: {log_path})" + raise RuntimeError(msg) + + data: dict[str, Any] = {} + for reader in output_readers: + data.update(reader() or {}) + + output_metadata: dict[str, Any] = {"returncode": result.returncode} + if include_stdio_metadata: + output_metadata.update({"stdout": result.stdout, "stderr": result.stderr}) + if metadata: + output_metadata.update(dict(metadata)) + + return StageOutput( + stage_name=stage_name, + data=data, + text=text_reader() if text_reader is not None else None, + logits=logits_reader() if logits_reader is not None else None, + timing_s=elapsed, + metadata=output_metadata, + ) + + +class HfTransformersReference: + """Run HuggingFace Transformers inference as the reference oracle.""" + + @property + def backend_name(self) -> str: + return "hf_transformers" + + def run_stage( + self, case: E2ECase, stage: StageSpec, ctx: RunContext + ) -> StageOutput: + if stage.name == "full_generation": + return self._run_full_generation(case, stage, ctx) + if stage.name == "full_inference": + return self._run_full_inference(case, stage, ctx) + if stage.name == "vision_encode": + # Vision encode is TRT-side only; reference skips this stage + return StageOutput( + stage_name=stage.name, + data={"skipped": True}, + metadata={"reason": "vision_encode handled by TRT runner only"}, + ) + raise ValueError(f"Unknown stage for hf_transformers: {stage.name!r}") + + def _run_full_generation( + self, case: E2ECase, stage: StageSpec, ctx: RunContext + ) -> StageOutput: + """Run HF model inference in a subprocess, collecting per-step logits. + + Dispatches to task-specific methods for non-standard tasks: + - vision_language_generation -> _run_vl_full_generation() + """ + task = case.task_strategy + if task == "vision_language_generation": + return self._run_vl_full_generation(case, stage, ctx) + + artifacts_dir = ctx.artifacts_dir or tempfile.gettempdir() + model_dir = _case_artifact_dir(artifacts_dir, case.name) if ctx.artifacts_dir else artifacts_dir + logits_path = str(Path(model_dir) / "hf_logits.npy") + text_path = str(Path(model_dir) / "hf_text.txt") + + prompt = case.inputs.get("prompt", "The capital of France is") + max_new_tokens = case.inputs.get("max_new_tokens", 30) + trust_remote_code = case.metadata.get("trust_remote_code", False) + hf_id = case.hf_id + model_ref = _resolve_cached_model_ref(hf_id) + torch_dtype_expr = _torch_dtype_for_case(case) + + contract_config = case.metadata.get("contract_config", {}) + use_chat_template = contract_config.get("use_chat_template", False) + enable_thinking = contract_config.get("enable_thinking", True) + + script = textwrap.dedent(f"""\ + import sys, numpy as np, torch + from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer + + hf_id = {hf_id!r} + model_ref = {model_ref!r} + prompt = {prompt!r} + max_new_tokens = {max_new_tokens} + trust_remote_code = {trust_remote_code!r} + logits_path = {logits_path!r} + text_path = {text_path!r} + use_chat_template = {use_chat_template!r} + enable_thinking = {enable_thinking!r} + + def _np(t): + return t.detach().float().cpu().numpy() + + tokenizer = AutoTokenizer.from_pretrained( + model_ref, trust_remote_code=trust_remote_code) + if use_chat_template: + messages = [{{"role": "user", "content": prompt}}] + try: + chat_kwargs = {{"add_generation_prompt": True}} + if not enable_thinking: + chat_kwargs["enable_thinking"] = False + text_input = tokenizer.apply_chat_template( + messages, tokenize=False, **chat_kwargs) + input_ids = tokenizer.encode(text_input, add_special_tokens=False) + except Exception: + # Fallback: model doesn't support chat template + input_ids = tokenizer.encode(prompt) + else: + input_ids = tokenizer.encode(prompt) + + load_kwargs = {{ + "trust_remote_code": trust_remote_code, + "torch_dtype": {torch_dtype_expr}, + }} + # Detect encoder-decoder models by checking config + from transformers import AutoConfig + _cfg = AutoConfig.from_pretrained(model_ref, trust_remote_code=trust_remote_code) + is_seq2seq = getattr(_cfg, "is_encoder_decoder", False) + + if is_seq2seq: + model = AutoModelForSeq2SeqLM.from_pretrained(model_ref, **load_kwargs) + else: + model = AutoModelForCausalLM.from_pretrained(model_ref, **load_kwargs) + model.eval() + + ids_tensor = torch.tensor([input_ids], dtype=torch.long) + all_logits = [] + + with torch.no_grad(): + if is_seq2seq: + # Encoder-decoder: use model.generate() for greedy decoding + output_ids = model.generate( + ids_tensor, max_new_tokens=max_new_tokens, + do_sample=False, num_beams=1) + generated_token_ids = output_ids[0].tolist() + # Re-run to get logits for each decoder step + decoder_ids = torch.tensor([generated_token_ids], dtype=torch.long) + outputs = model(ids_tensor, decoder_input_ids=decoder_ids) + for i in range(outputs.logits.shape[1]): + all_logits.append(_np(outputs.logits[0, i])) + text = tokenizer.decode(generated_token_ids, skip_special_tokens=True) + else: + # Decoder-only: step-by-step autoregressive + outputs = model(ids_tensor) + prefill_logits = _np(outputs.logits[0]) + for i in range(len(input_ids)): + all_logits.append(prefill_logits[i]) + + gen_ids = list(input_ids) + generated_token_ids = [] + eos_id = getattr(tokenizer, "eos_token_id", None) + for _ in range(max_new_tokens): + next_token = int(np.argmax(all_logits[-1])) + generated_token_ids.append(next_token) + if eos_id is not None and next_token == eos_id: + break + gen_ids.append(next_token) + ids_tensor = torch.tensor([gen_ids], dtype=torch.long) + outputs = model(ids_tensor) + all_logits.append(_np(outputs.logits[0, -1])) + text = tokenizer.decode(generated_token_ids, skip_special_tokens=True) + + with open(text_path, "w") as f: + f.write(text) + + # Pad and save logits + max_len = max(l.shape[0] for l in all_logits) + padded = np.zeros((len(all_logits), max_len), dtype=np.float32) + for i, l in enumerate(all_logits): + padded[i, :l.shape[0]] = l + np.save(logits_path, padded) + + print(f"OK steps={{len(all_logits)}} vocab={{max_len}}") + """) + + python = ctx.reference_python_path() or sys.executable + logger.info("HF reference: running %s", case.name) + return run_reference_subprocess( + command=[python, "-c", script], + timeout_s=1800, + label="hf_full_generation", + artifact_dir=ctx.artifacts_dir or "", + case_name=case.name, + stage_name=stage.name, + env=_reference_env(ctx), + output_readers=(_existing_path_reader(logits_path, "logits_path"),), + text_reader=lambda: _read_text_artifact(text_path), + logits_reader=( + lambda: logits_path if Path(logits_path).is_file() else None + ), + metadata={"trust_remote_code": trust_remote_code}, + include_stdio_metadata=True, + failure_label="HF reference", + ) + + def _run_full_inference( + self, case: E2ECase, stage: StageSpec, ctx: RunContext + ) -> StageOutput: + """Run HF model forward pass for non-generative tasks. + + Dispatches based on task_strategy to the appropriate HF Auto class. + """ + task = case.task_strategy + if task == "encoder_only_nlp": + return self._run_encoder_only(case, stage, ctx) + if task == "segmentation": + return self._run_segmentation_ref(case, stage, ctx) + if task == "embedding": + return self._run_embedding_ref(case, stage, ctx) + if task == "reranking": + return self._run_reranking_ref(case, stage, ctx) + if task == "object_detection": + return self._run_object_detection_ref(case, stage, ctx) + if task == "image_classification": + return self._run_image_classification_ref(case, stage, ctx) + raise ValueError( + f"full_inference not implemented for task_strategy={task!r}") + + def _run_encoder_only( + self, case: E2ECase, stage: StageSpec, ctx: RunContext + ) -> StageOutput: + """Run HF encoder-only model (e.g. BERT) and return CLS embedding.""" + artifacts_dir = ctx.artifacts_dir or tempfile.gettempdir() + model_dir = _case_artifact_dir(artifacts_dir, case.name) if ctx.artifacts_dir else artifacts_dir + output_path = str(Path(model_dir) / "hf_encoder.json") + + prompt = case.inputs.get("prompt", "Hello world") + trust_remote_code = case.metadata.get("trust_remote_code", False) + hf_id = case.hf_id + model_ref = _resolve_cached_model_ref(hf_id) + torch_dtype_expr = _torch_dtype_for_case(case) + + script = textwrap.dedent(f"""\ + import json, torch, numpy as np + from transformers import AutoModel, AutoTokenizer + + hf_id = {hf_id!r} + model_ref = {model_ref!r} + prompt = {prompt!r} + trust_remote_code = {trust_remote_code!r} + output_path = {output_path!r} + + tokenizer = AutoTokenizer.from_pretrained( + model_ref, trust_remote_code=trust_remote_code) + model = AutoModel.from_pretrained( + model_ref, trust_remote_code=trust_remote_code, + torch_dtype={torch_dtype_expr}) + model.eval() + + inputs = tokenizer(prompt, return_tensors="pt") + with torch.no_grad(): + outputs = model(**inputs) + + # CLS token embedding from last_hidden_state + if hasattr(outputs, 'last_hidden_state') and outputs.last_hidden_state is not None: + cls_embedding = outputs.last_hidden_state[0, 0].float().cpu().numpy().tolist() + else: + first_out = outputs[0] + if first_out.ndim == 3: + cls_embedding = first_out[0, 0].float().cpu().numpy().tolist() + elif first_out.ndim == 2: + cls_embedding = first_out[0].float().cpu().numpy().tolist() + else: + cls_embedding = first_out.float().cpu().numpy().tolist() + result = {{"cls_embedding": cls_embedding}} + with open(output_path, "w") as f: + json.dump(result, f) + print("OK") + """) + + python = ctx.reference_python_path() or sys.executable + return run_reference_subprocess( + command=[python, "-c", script], + timeout_s=600, + label="hf_encoder_only", + artifact_dir=ctx.artifacts_dir or "", + case_name=case.name, + stage_name=stage.name, + env=_reference_env(ctx), + output_readers=(_json_output_reader(output_path),), + failure_label="HF encoder-only", + ) + + @staticmethod + def _resolve_image_path(image_path: str) -> str: + """Resolve image path, handling relative paths from manifests.""" + if not image_path: + return image_path + if os.path.isabs(image_path): + return image_path + # Resolve relative to tests/e2e/ directory + resolved = E2E_DIR / image_path + if resolved.exists(): + return str(resolved) + # Also try relative to project root + resolved2 = PROJECT_DIR / image_path + if resolved2.exists(): + return str(resolved2) + return image_path + + def _run_embedding_ref( + self, case: E2ECase, stage: StageSpec, ctx: RunContext + ) -> StageOutput: + """Run HF embedding model as reference — mean pool + L2 normalize.""" + artifacts_dir = ctx.artifacts_dir or tempfile.gettempdir() + model_dir = _case_artifact_dir(artifacts_dir, case.name) if ctx.artifacts_dir else artifacts_dir + output_path = str(Path(model_dir) / "hf_embedding.json") + + prompt = case.inputs.get("prompt", "What is machine learning?") + trust_remote_code = case.metadata.get("trust_remote_code", False) + hf_id = case.hf_id + model_ref = _resolve_cached_model_ref(hf_id) + torch_dtype_expr = _torch_dtype_for_case(case) + + script = textwrap.dedent(f"""\ + import json, torch, numpy as np + from transformers import AutoModel, AutoTokenizer + + hf_id = {hf_id!r} + model_ref = {model_ref!r} + prompt = {prompt!r} + trust_remote_code = {trust_remote_code!r} + output_path = {output_path!r} + + tokenizer = AutoTokenizer.from_pretrained( + model_ref, trust_remote_code=trust_remote_code) + model = AutoModel.from_pretrained( + model_ref, trust_remote_code=trust_remote_code, + torch_dtype={torch_dtype_expr}) + model.eval() + + # Generic forward pass: tokenize -> forward -> mean pool -> L2 norm + # (We use the raw forward pass to match TRT, not encode_queries() + # which adds an instruction prefix that TRT doesn't replicate.) + inputs = tokenizer(prompt, return_tensors="pt", padding=True, + truncation=True) + with torch.no_grad(): + outputs = model(**inputs, output_hidden_states=True) + # Try last_hidden_state first, then fall back to hidden_states[-1] + if hasattr(outputs, "last_hidden_state") and outputs.last_hidden_state is not None: + hidden = outputs.last_hidden_state + elif hasattr(outputs, "hidden_states") and outputs.hidden_states: + hidden = outputs.hidden_states[-1] + else: + raise RuntimeError("Model output has no hidden states") + mask = inputs["attention_mask"].unsqueeze(-1).float() + pooled = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9) + pooled = torch.nn.functional.normalize(pooled, p=2, dim=-1) + embedding = pooled[0].float().cpu().numpy().tolist() + + result = {{"embedding": embedding}} + with open(output_path, "w") as f: + json.dump(result, f) + print("OK") + """) + + python = ctx.reference_python_path() or sys.executable + return run_reference_subprocess( + command=[python, "-c", script], + timeout_s=600, + label="hf_embedding", + artifact_dir=ctx.artifacts_dir or "", + case_name=case.name, + stage_name=stage.name, + env=_reference_env(ctx), + output_readers=(_json_output_reader(output_path),), + failure_label="HF embedding ref", + ) + + def _run_image_classification_ref( + self, case: E2ECase, stage: StageSpec, ctx: RunContext + ) -> StageOutput: + """Run timm image classification as the reference oracle.""" + artifacts_dir = ctx.artifacts_dir or tempfile.gettempdir() + model_dir = _case_artifact_dir(artifacts_dir, case.name) if ctx.artifacts_dir else artifacts_dir + output_path = str(Path(model_dir) / "hf_image_classification.json") + + image_path = self._resolve_image_path(case.inputs.get("image", "")) + hf_id = case.hf_id + + script = textwrap.dedent(f"""\ + import json + from pathlib import Path + + import numpy as np + import timm + import torch + from PIL import Image + from timm.data import create_transform, resolve_model_data_config + + hf_id = {hf_id!r} + image_path = {image_path!r} + output_path = {output_path!r} + + model_ref = f"hf-hub:{{hf_id}}" + try: + model = timm.create_model(model_ref, pretrained=True) + except Exception: + model = timm.create_model(f"hf_hub:{{hf_id}}", pretrained=True) + model.eval() + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model.to(device) + data_config = resolve_model_data_config(model) + transform = create_transform(**data_config, is_training=False) + image = Image.open(image_path).convert("RGB") + tensor = transform(image).unsqueeze(0).to(device) + + with torch.no_grad(): + logits = model(tensor)[0].float().cpu().numpy() + + top_class = int(np.argmax(logits)) + result = {{ + "top_class": top_class, + "top_score": float(logits[top_class]), + "num_classes": int(logits.shape[0]), + }} + Path(output_path).parent.mkdir(parents=True, exist_ok=True) + with open(output_path, "w") as f: + json.dump(result, f) + print("OK") + """) + + python = ctx.reference_python_path() or sys.executable + return run_reference_subprocess( + command=[python, "-c", script], + timeout_s=900, + label="hf_image_classification", + artifact_dir=ctx.artifacts_dir or "", + case_name=case.name, + stage_name=stage.name, + env=_reference_env(ctx), + output_readers=(_json_output_reader(output_path),), + include_stdio_metadata=True, + failure_label="HF image classification", + ) + + def _run_segmentation_ref( + self, case: E2ECase, stage: StageSpec, ctx: RunContext + ) -> StageOutput: + """Run HF segmentation model as reference.""" + artifacts_dir = ctx.artifacts_dir or tempfile.gettempdir() + model_dir = _case_artifact_dir(artifacts_dir, case.name) if ctx.artifacts_dir else artifacts_dir + output_path = str(Path(model_dir) / "hf_seg.npy") + + image_path = self._resolve_image_path(case.inputs.get("image", "")) + trust_remote_code = case.metadata.get("trust_remote_code", False) + hf_id = case.hf_id + torch_dtype_expr = _torch_dtype_for_case(case) + + script = textwrap.dedent(f"""\ + import numpy as np, torch + from transformers import AutoModelForSemanticSegmentation, AutoImageProcessor + from PIL import Image + + hf_id = {hf_id!r} + image_path = {image_path!r} + trust_remote_code = {trust_remote_code!r} + output_path = {output_path!r} + + processor = AutoImageProcessor.from_pretrained( + hf_id, trust_remote_code=trust_remote_code) + model = AutoModelForSemanticSegmentation.from_pretrained( + hf_id, trust_remote_code=trust_remote_code, + torch_dtype={torch_dtype_expr}) + model.eval() + + image = Image.open(image_path).convert("RGB") + inputs = processor(images=image, return_tensors="pt") + with torch.no_grad(): + outputs = model(**inputs) + logits = outputs.logits[0].float().cpu().numpy() + class_map = np.argmax(logits, axis=0).astype(np.int32) + np.save(output_path, class_map) + print(f"OK classes={{class_map.max() + 1}}") + """) + + def _segmentation_outputs() -> dict[str, Any]: + data: dict[str, Any] = {} + if Path(output_path).is_file(): + data["class_map_path"] = output_path + import numpy as np + + data["class_map"] = np.load(output_path) + + try: + from PIL import Image + + cmap = data["class_map"] + num_classes = int(cmap.max()) + 1 + np.random.seed(42) + palette = np.random.randint( + 0, 255, (num_classes, 3), dtype=np.uint8 + ) + palette[0] = [0, 0, 0] + colored = palette[cmap] + viz_path = output_path.replace(".npy", "_viz.png") + Image.fromarray(colored).save(viz_path) + data["viz_path"] = viz_path + except Exception as e: + logger.warning("Failed to save segmentation viz: %s", e) + return data + + python = ctx.reference_python_path() or sys.executable + return run_reference_subprocess( + command=[python, "-c", script], + timeout_s=600, + label="hf_segmentation", + artifact_dir=ctx.artifacts_dir or "", + case_name=case.name, + stage_name=stage.name, + env=_reference_env(ctx), + output_readers=(_segmentation_outputs,), + failure_label="HF segmentation", + ) + + def _run_reranking_ref( + self, case: E2ECase, stage: StageSpec, ctx: RunContext + ) -> StageOutput: + """Run HF cross-encoder reranking and return one score per document.""" + artifacts_dir = ctx.artifacts_dir or tempfile.gettempdir() + model_dir = _case_artifact_dir(artifacts_dir, case.name) if ctx.artifacts_dir else artifacts_dir + output_path = str(Path(model_dir) / "hf_rerank.json") + + prompt = case.inputs.get("prompt", "query: test") + documents = case.inputs.get("documents") + if documents is None: + document = case.inputs.get("document", "") + documents = [document] if document else [] + trust_remote_code = case.metadata.get("trust_remote_code", False) + hf_id = case.hf_id + model_ref = _resolve_cached_model_ref(hf_id) + torch_dtype_expr = _torch_dtype_for_case(case) + + script = textwrap.dedent(f"""\ + import json, torch + from transformers import AutoModelForSequenceClassification, AutoProcessor, AutoTokenizer + + hf_id = {hf_id!r} + model_ref = {model_ref!r} + prompt = {prompt!r} + documents = {documents!r} + trust_remote_code = {trust_remote_code!r} + output_path = {output_path!r} + torch_dtype = {torch_dtype_expr} + + model = AutoModelForSequenceClassification.from_pretrained( + model_ref, trust_remote_code=trust_remote_code, + torch_dtype=torch_dtype) + device = "cuda" if torch.cuda.is_available() else "cpu" + model.to(device) + model.eval() + + examples = [ + {{"question": prompt, "doc_text": doc, "doc_image": ""}} + for doc in documents + ] + + try: + processor = AutoProcessor.from_pretrained( + model_ref, trust_remote_code=trust_remote_code, + max_input_tiles=6, use_thumbnail=True, + rerank_max_length=8192) + if not hasattr(processor, "process_queries_documents_crossencoder"): + raise AttributeError("processor has no cross-encoder helper") + inputs = processor.process_queries_documents_crossencoder(examples) + except Exception: + tokenizer = AutoTokenizer.from_pretrained( + model_ref, trust_remote_code=trust_remote_code) + texts = [ + f"question:{{prompt}} passage:{{doc}}" + for doc in documents + ] + inputs = tokenizer( + texts, return_tensors="pt", padding=True, truncation=True) + + inputs = {{ + key: value.to(device) if hasattr(value, "to") else value + for key, value in inputs.items() + }} + with torch.no_grad(): + outputs = model(**inputs) + logits = outputs.logits.detach().float().cpu() + if logits.ndim == 2 and logits.shape[-1] == 1: + scores = logits[:, 0].tolist() + elif logits.ndim == 2: + scores = logits[:, -1].tolist() + else: + scores = logits.reshape(-1).tolist() + result = {{"scores": scores}} + with open(output_path, "w") as f: + json.dump(result, f) + print("OK") + """) + + python = ctx.reference_python_path() or sys.executable + return run_reference_subprocess( + command=[python, "-c", script], + timeout_s=600, + label="hf_reranking", + artifact_dir=ctx.artifacts_dir or "", + case_name=case.name, + stage_name=stage.name, + env=_reference_env(ctx), + output_readers=(_json_output_reader(output_path),), + failure_label="HF reranking", + ) + + def _run_object_detection_ref( + self, case: E2ECase, stage: StageSpec, ctx: RunContext + ) -> StageOutput: + """Run HF object detection model as reference.""" + artifacts_dir = ctx.artifacts_dir or tempfile.gettempdir() + model_dir = _case_artifact_dir(artifacts_dir, case.name) if ctx.artifacts_dir else artifacts_dir + output_path = str(Path(model_dir) / "hf_det.json") + + image_path = self._resolve_image_path(case.inputs.get("image", "")) + trust_remote_code = case.metadata.get("trust_remote_code", False) + hf_id = case.hf_id + torch_dtype_expr = _torch_dtype_for_case(case) + + script = textwrap.dedent(f"""\ + import json, torch + from transformers import AutoModelForObjectDetection, AutoImageProcessor + from PIL import Image + + hf_id = {hf_id!r} + image_path = {image_path!r} + trust_remote_code = {trust_remote_code!r} + output_path = {output_path!r} + + processor = AutoImageProcessor.from_pretrained( + hf_id, trust_remote_code=trust_remote_code) + model = AutoModelForObjectDetection.from_pretrained( + hf_id, trust_remote_code=trust_remote_code, + torch_dtype={torch_dtype_expr}) + model.eval() + + image = Image.open(image_path).convert("RGB") + inputs = processor(images=image, return_tensors="pt") + with torch.no_grad(): + outputs = model(**inputs) + # Post-process to get boxes + scores + target_sizes = torch.tensor([image.size[::-1]]) + results = processor.post_process_object_detection( + outputs, target_sizes=target_sizes, threshold=0.5)[0] + detections = [] + for score, label, box in zip( + results["scores"], results["labels"], results["boxes"] + ): + detections.append({{ + "score": score.item(), + "label": label.item(), + "box": box.tolist(), + }}) + with open(output_path, "w") as f: + json.dump({{"detections": detections}}, f) + print(f"OK detections={{len(detections)}}") + """) + + python = ctx.reference_python_path() or sys.executable + return run_reference_subprocess( + command=[python, "-c", script], + timeout_s=600, + label="hf_object_detection", + artifact_dir=ctx.artifacts_dir or "", + case_name=case.name, + stage_name=stage.name, + env=_reference_env(ctx), + output_readers=(_json_output_reader(output_path),), + failure_label="HF object detection", + ) + + def _run_vl_full_generation( + self, case: E2ECase, stage: StageSpec, ctx: RunContext + ) -> StageOutput: + """Run HF vision-language model for reference generation.""" + artifacts_dir = ctx.artifacts_dir or tempfile.gettempdir() + model_dir = _case_artifact_dir(artifacts_dir, case.name) if ctx.artifacts_dir else artifacts_dir + text_path = str(Path(model_dir) / "hf_vl_text.txt") + + prompt = case.inputs.get("prompt", "Describe this image.") + max_new_tokens = case.inputs.get("max_new_tokens", 30) + trust_remote_code = case.metadata.get("trust_remote_code", False) + image_path = self._resolve_image_path(case.inputs.get("image", "")) + hf_id = case.hf_id + model_ref = _resolve_cached_model_ref(hf_id) + fallback_text = prompt + torch_dtype_expr = _torch_dtype_for_case(case) + + script = textwrap.dedent(f"""\ + import sys, torch + from transformers import AutoProcessor + from PIL import Image + from {__name__} import ( + _decode_vl_generated_text, + _vl_prompt_has_image_placeholder, + ) + + hf_id = {hf_id!r} + model_ref = {model_ref!r} + prompt = {prompt!r} + fallback_text = {fallback_text!r} + max_new_tokens = {max_new_tokens} + trust_remote_code = {trust_remote_code!r} + image_path = {image_path!r} + text_path = {text_path!r} + + processor = AutoProcessor.from_pretrained( + model_ref, trust_remote_code=trust_remote_code) + + # Try VL-specific auto classes in preference order + import transformers + model = None + for cls_name in ["AutoModelForImageTextToText", + "AutoModelForVision2Seq"]: + try: + cls = getattr(transformers, cls_name) + model = cls.from_pretrained( + model_ref, trust_remote_code=trust_remote_code, + torch_dtype={torch_dtype_expr}) + break + except (AttributeError, ImportError, ValueError, KeyError): + continue + # Fallback for models registered as causal LM with multimodal + # inputs (e.g. Phi-4-multimodal) + if model is None: + model = transformers.AutoModelForCausalLM.from_pretrained( + model_ref, trust_remote_code=True, + torch_dtype={torch_dtype_expr}) + model.eval() + + image = Image.open(image_path).convert("RGB") + + # Build conversation for chat-template models + messages = [ + {{"role": "user", "content": [ + {{"type": "image", "image": image_path}}, + {{"type": "text", "text": prompt}}, + ]}} + ] + text_input = "" + try: + text_input = processor.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True) + if not isinstance(text_input, str): + raise TypeError("processor.apply_chat_template did not return text") + if not _vl_prompt_has_image_placeholder(text_input): + raise ValueError("chat template produced no image placeholder") + inputs = processor( + text=text_input, images=image, return_tensors="pt") + except Exception: + # Fallback for models without chat template + inputs = processor( + text=fallback_text, images=image, return_tensors="pt") + + with torch.no_grad(): + generated_ids = model.generate( + **inputs, max_new_tokens=max_new_tokens) + + # Decode only the generated portion (after input) + input_len = inputs.get("input_ids", torch.tensor([])).shape[-1] + text = _decode_vl_generated_text( + processor, + generated_ids[0], + input_len, + (prompt, fallback_text, text_input), + ) + if not text.strip(): + raise RuntimeError( + "HF VL reference produced empty or prompt-only generated text") + + with open(text_path, "w") as f: + f.write(text) + print(f"OK text={{text[:100]!r}}") + """) + + python = ctx.reference_python_path() or sys.executable + return run_reference_subprocess( + command=[python, "-c", script], + timeout_s=1800, + label="hf_vl_generation", + artifact_dir=ctx.artifacts_dir or "", + case_name=case.name, + stage_name=stage.name, + env=_reference_env(ctx), + output_readers=(lambda: {"text": _read_text_artifact(text_path)},), + text_reader=lambda: _read_text_artifact(text_path), + metadata={"trust_remote_code": trust_remote_code}, + failure_label="HF VL generation", + ) + + +plugin = HfTransformersReference() diff --git a/tests/e2e/models/timm_repvgg/e2e_plugins/references/invariant_only.py b/tests/e2e/models/timm_repvgg/e2e_plugins/references/invariant_only.py new file mode 100644 index 0000000000..b99cc47bfb --- /dev/null +++ b/tests/e2e/models/timm_repvgg/e2e_plugins/references/invariant_only.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Invariant-only reference backend — no external reference needed. + +Returns a dummy StageOutput. The comparator only checks invariants +(nan, range, shape) on the TRT output itself. Used for models where +no reference implementation is available. +""" + +from __future__ import annotations + +import logging + +from ..contracts import E2ECase, RunContext, StageOutput, StageSpec + +logger = logging.getLogger(__name__) + + +class InvariantOnlyReference: + """No-op reference backend for invariant-only testing.""" + + @property + def backend_name(self) -> str: + return "invariant_only" + + def run_stage( + self, case: E2ECase, stage: StageSpec, ctx: RunContext + ) -> StageOutput: + logger.info( + "Invariant-only reference for %s/%s — returning dummy output", + case.name, stage.name, + ) + return StageOutput( + stage_name=stage.name, + data={"_invariant_only": True}, + text=None, + timing_s=0.0, + metadata={ + "source": "invariant_only", + "note": "No external reference; comparator checks invariants only", + }, + ) + + +plugin = InvariantOnlyReference() diff --git a/tests/e2e/models/timm_repvgg/e2e_plugins/references/nemo_reference.py b/tests/e2e/models/timm_repvgg/e2e_plugins/references/nemo_reference.py new file mode 100644 index 0000000000..1b992cf533 --- /dev/null +++ b/tests/e2e/models/timm_repvgg/e2e_plugins/references/nemo_reference.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Generic NeMo reference backend registration. + +Model-specific NeMo reference implementations live under +``tests/e2e/models//e2e_plugins/references``. +""" + +from __future__ import annotations + +from ..contracts import E2ECase, RunContext, StageOutput, StageSpec + + +class NemoReference: + """Reference backend placeholder for model-owned NeMo implementations.""" + + @property + def backend_name(self) -> str: + return "nemo" + + def run_stage( + self, case: E2ECase, stage: StageSpec, ctx: RunContext + ) -> StageOutput: + raise ValueError( + "Shared NeMo reference backend does not implement model-specific " + f"task_strategy={case.task_strategy!r} " + f"runtime_strategy={case.runtime_strategy!r}" + ) + + +plugin = NemoReference() diff --git a/tests/e2e/models/timm_repvgg/e2e_plugins/registry.py b/tests/e2e/models/timm_repvgg/e2e_plugins/registry.py new file mode 100644 index 0000000000..69e6bba90b --- /dev/null +++ b/tests/e2e/models/timm_repvgg/e2e_plugins/registry.py @@ -0,0 +1,10 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Model-local bridge to the active E2E registry.""" + +from tests.e2e_harness.registry import ( # noqa: F401 + register_comparator, + register_reference, + register_runner, +) diff --git a/tests/e2e/models/timm_repvgg/e2e_plugins/repro.py b/tests/e2e/models/timm_repvgg/e2e_plugins/repro.py new file mode 100644 index 0000000000..a69e20459a --- /dev/null +++ b/tests/e2e/models/timm_repvgg/e2e_plugins/repro.py @@ -0,0 +1,52 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""timm_repvgg model-owned E2E repro command provider.""" + +from __future__ import annotations + +import shlex + +from .contracts import E2ECase, ReproCommandProvider, RunContext + + +def _shell_quote(value: object) -> str: + return shlex.quote(str(value)) + + +class TimmRepvggReproCommandProvider: + """Build timm_repvgg TRT repro commands without shared harness branches.""" + + @property + def family_name(self) -> str: + return "timm_repvgg" + + def build_trt_inference_command( + self, + case: E2ECase, + ctx: RunContext, + bundle_path: str, + ) -> list[str] | None: + if case.task_strategy != "image_classification": + return None + + image = ( + case.inputs.get("image") + or case.inputs.get("test_image") + or case.inputs.get("image_path") + or "" + ) + infer_parts = [ + ctx.binary_path, + "classify", + bundle_path, + "--image", + _shell_quote(image), + ] + runtime_cli_python = ctx.runtime_cli_hf_python() + if runtime_cli_python: + infer_parts.extend(["--hf-python", runtime_cli_python]) + return infer_parts + + +repro_provider: ReproCommandProvider = TimmRepvggReproCommandProvider() diff --git a/tests/e2e/models/timm_repvgg/e2e_plugins/runner.py b/tests/e2e/models/timm_repvgg/e2e_plugins/runner.py new file mode 100644 index 0000000000..096d4cfbea --- /dev/null +++ b/tests/e2e/models/timm_repvgg/e2e_plugins/runner.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""timm_repvgg model-owned E2E runner plugins.""" + +from __future__ import annotations + +from .runners.image_classification import ImageClassificationRunner + + +class TimmRepvggImageClassificationRunner(ImageClassificationRunner): + """timm_repvgg local runner for image_classification.""" + +runner = TimmRepvggImageClassificationRunner() diff --git a/tests/e2e/models/timm_repvgg/e2e_plugins/runners/__init__.py b/tests/e2e/models/timm_repvgg/e2e_plugins/runners/__init__.py new file mode 100644 index 0000000000..986743e50e --- /dev/null +++ b/tests/e2e/models/timm_repvgg/e2e_plugins/runners/__init__.py @@ -0,0 +1,9 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Strategy runners — TRT inference execution for each task strategy. + +Each module in this package should expose a module-level ``plugin`` attribute +that is an instance implementing the TaskStrategyRunner protocol. The registry +auto-discovers these plugins on first access. +""" diff --git a/tests/e2e/models/timm_repvgg/e2e_plugins/runners/_runtime_common.py b/tests/e2e/models/timm_repvgg/e2e_plugins/runners/_runtime_common.py new file mode 100644 index 0000000000..7d4cb2b453 --- /dev/null +++ b/tests/e2e/models/timm_repvgg/e2e_plugins/runners/_runtime_common.py @@ -0,0 +1,312 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Family-local E2E subprocess and distributed runtime helpers.""" + +from __future__ import annotations + +import logging +import json +import re +import shutil +import subprocess +import tempfile +from pathlib import Path + +from .. import _case_artifact_dir +from ..contracts import E2ECase, RunContext + +logger = logging.getLogger(__name__) + +_SUPPORTED_STAGES = {"full_generation", "prefill", "decode"} +_TRTMC_TIMING_RE = re.compile( + r"^\[trtmc\.timing\]\s+" + r"prefill_ms=(?P[-+0-9.eE]+)\s+" + r"decode_ms=(?P[-+0-9.eE]+)\s+" + r"total_ms=(?P[-+0-9.eE]+)\s*$", + re.MULTILINE, +) +_TRTMC_LOAD_TIMING_RE = re.compile( + r"^\[trtmc\.load_timing\]\s+.*?" + r"load_deserialize_ms=(?P[-+0-9.eE]+)", + re.MULTILINE, +) +_TRT_RUNTIME_ERROR_RE = re.compile( + r"(?im)^.*(" + r"\[trt\]\s+ERROR:" + r"|IExecutionContext::enqueueV3:\s+Error Code" + r"|Internal Error:" + r"|Cuda Runtime" + r"|illegal memory access" + r").*$" +) +_MPI_TAGGED_STDOUT_RE = re.compile( + r"^\[[^\]]+,(?P\d+)\]:(?P.*)$") +_MPI_STREAM_TAG_RE = re.compile(r"\[[^\]]+,\d+\]<(?:stdout|stderr)>:") + + +def _distributed_runtime_config(case: E2ECase | None) -> dict: + if case is None: + return {} + config = case.metadata.get("distributed_runtime", {}) + return config if isinstance(config, dict) and config.get("enabled") else {} + + +def _extract_rank_zero_stdout(stdout: str) -> str: + """Return rank-0 stdout from OpenMPI --tag-output, falling back to raw text.""" + rank0_lines: list[str] = [] + saw_tagged = False + for line in (stdout or "").splitlines(): + match = _MPI_TAGGED_STDOUT_RE.match(line) + if match is None: + continue + saw_tagged = True + if int(match.group("rank")) == 0: + rank0_lines.append(match.group("text")) + if saw_tagged: + return "\n".join(rank0_lines).strip() + return (stdout or "").strip() + + +def _strip_mpi_stream_tags(text: str) -> str: + return _MPI_STREAM_TAG_RE.sub("", text or "") + + +def _safe_artifact_name(name: str) -> str: + return re.sub(r"[^A-Za-z0-9_.-]+", "_", name or "case") + + +def _read_text_generation_sample(path: Path) -> dict: + """Read the first JSONL text-generation sample written by the C++ CLI.""" + if not path.is_file(): + return {} + with path.open("r", encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line: + continue + sample = json.loads(line) + if not isinstance(sample, dict): + return {} + token_ids = sample.get("token_ids") + if isinstance(token_ids, list): + sample["token_ids"] = [int(token) for token in token_ids] + return sample + return {} + + +def _ensure_distributed_runtime_env( + case: E2ECase, + ctx: RunContext, + env: dict[str, str], + rendezvous_suffix: str = "", +) -> None: + """Populate shared env values needed by all distributed ranks.""" + if not _distributed_runtime_config(case): + return + if env.get("TRTMC_NCCL_RENDEZVOUS"): + return + + safe_name = _safe_artifact_name(case.name) + root = Path(_case_artifact_dir(ctx.artifacts_dir, case.name)) if ctx.artifacts_dir else \ + Path(tempfile.gettempdir()) + path = root / f"{safe_name}{rendezvous_suffix}.nccl_rendezvous.bin" + path.parent.mkdir(parents=True, exist_ok=True) + try: + path.unlink() + except FileNotFoundError: + pass + env["TRTMC_NCCL_RENDEZVOUS"] = str(path) + + +def _wrap_distributed_command( + cmd: list[str], case: E2ECase | None, env: dict[str, str] +) -> list[str]: + config = _distributed_runtime_config(case) + if not config: + return cmd + + launcher = str(config.get("launcher", "mpirun") or "mpirun") + world_size = int(config.get("world_size", config.get("tp_size", 2)) or 2) + launcher_args = config.get("launcher_args") + if isinstance(launcher_args, list): + prefix = [launcher] + [str(arg) for arg in launcher_args] + else: + prefix = [launcher, "--tag-output", "-np", str(world_size)] + + export_env = config.get("export_env", ["LD_LIBRARY_PATH", "CUDA_VISIBLE_DEVICES"]) + if isinstance(export_env, list) and Path(launcher).name == "mpirun": + export_names = [str(name) for name in export_env] + for name in ("TRTMC_NCCL_RENDEZVOUS", "TRTMC_EMBEDDING_STDOUT"): + if name in env and name not in export_names: + export_names.append(name) + for name in export_names: + if name in env: + prefix.extend(["-x", name]) + + return prefix + cmd + + +def _visible_gpu_indices(env: dict[str, str]) -> list[str]: + raw = env.get("CUDA_VISIBLE_DEVICES", "") + if not raw or raw.lower() in {"all", "none", "void"}: + return [] + indices: list[str] = [] + for part in raw.split(","): + token = part.strip() + if token.isdigit(): + indices.append(token) + return indices + + +class _GpuMemorySampler: + def __init__(self, artifacts_dir: str | None, case_name: str, env: dict[str, str], + interval_ms: int) -> None: + root = Path(_case_artifact_dir(artifacts_dir, case_name)) if artifacts_dir else \ + Path(tempfile.gettempdir()) + root.mkdir(parents=True, exist_ok=True) + self.path = root / "gpu_memory_samples.csv" + self.env = env + self.interval_ms = max(50, interval_ms) + self.visible_indices = _visible_gpu_indices(env) + self.proc: subprocess.Popen | None = None + self.handle = None + self.error = "" + + def start(self) -> None: + if shutil.which("nvidia-smi") is None: + self.error = "nvidia-smi not found" + return + self.handle = self.path.open("w", encoding="utf-8") + cmd = [ + "nvidia-smi", + "--query-gpu=index,memory.used", + "--format=csv,noheader,nounits", + f"--loop-ms={self.interval_ms}", + ] + try: + self.proc = subprocess.Popen( + cmd, + stdout=self.handle, + stderr=subprocess.DEVNULL, + text=True, + env=self.env, + ) + except Exception as exc: + self.error = str(exc) + self.handle.close() + self.handle = None + + def stop(self) -> dict: + if self.proc is not None: + self.proc.terminate() + try: + self.proc.wait(timeout=2) + except subprocess.TimeoutExpired: + self.proc.kill() + self.proc.wait(timeout=2) + if self.handle is not None: + self.handle.close() + self.handle = None + return self._summary() + + def _summary(self) -> dict: + meta = { + "sample_file": str(self.path), + "sample_interval_ms": self.interval_ms, + "visible_device_indices": self.visible_indices, + } + if self.error: + meta["error"] = self.error + return meta + peaks: dict[str, int] = {} + sample_count = 0 + if not self.path.is_file(): + meta["error"] = "sample file was not created" + return meta + with self.path.open("r", encoding="utf-8") as f: + for line in f: + parts = [p.strip() for p in line.split(",")] + if len(parts) < 2 or not parts[0].isdigit(): + continue + if self.visible_indices and parts[0] not in self.visible_indices: + continue + try: + used_mb = int(float(parts[1])) + except ValueError: + continue + peaks[parts[0]] = max(peaks.get(parts[0], 0), used_mb) + sample_count += 1 + meta["sample_count"] = sample_count + meta["peak_memory_mb_by_gpu"] = peaks + if peaks: + meta["peak_memory_mb"] = max(peaks.values()) + meta["peak_memory_mb_visible_sum"] = sum(peaks.values()) + return meta + + +def _maybe_start_gpu_memory_sampler( + distributed_runtime: dict, ctx: RunContext, case: E2ECase | None, env: dict[str, str] +) -> _GpuMemorySampler | None: + if case is None or not distributed_runtime.get("capture_gpu_memory"): + return None + interval_ms = int(distributed_runtime.get("gpu_memory_sample_interval_ms", 200) or 200) + sampler = _GpuMemorySampler(ctx.artifacts_dir, case.name, env, interval_ms) + sampler.start() + return sampler + + +def _extract_trtmc_timing(stderr: str) -> dict[str, float]: + match = _TRTMC_TIMING_RE.search(stderr or "") + if match is None: + return {} + try: + prefill_ms = float(match.group("prefill_ms")) + decode_ms = float(match.group("decode_ms")) + total_ms = float(match.group("total_ms")) + except ValueError: + return {} + return { + "trt_engine_prefill_s": prefill_ms / 1000.0, + "trt_engine_decode_s": decode_ms / 1000.0, + "trt_engine_s": total_ms / 1000.0, + } + + +def _extract_trtmc_load_timing(stderr: str) -> dict[str, float]: + total_ms = 0.0 + found = False + for match in _TRTMC_LOAD_TIMING_RE.finditer(stderr or ""): + try: + total_ms += float(match.group("load_deserialize_ms")) + found = True + except ValueError: + continue + return {"trt_load_deserialize_s": total_ms / 1000.0} if found else {} + + +def _detect_trt_runtime_error(stderr: str) -> str: + match = _TRT_RUNTIME_ERROR_RE.search(stderr or "") + return match.group(0).strip() if match else "" + + +def _distributed_debug_logits_required(case: E2ECase) -> bool: + distributed_runtime = _distributed_runtime_config(case) + return bool(distributed_runtime and distributed_runtime.get("debug_logits", True)) + + +def _format_debug_runner_error(case: E2ECase, phase: str, meta: dict) -> str: + detail = meta.get("error") + if not detail and meta.get("returncode") not in (None, 0): + detail = f"returncode={meta['returncode']}" + if not detail: + detail = "logits were not produced" + log_path = meta.get("stderr_log") + if log_path: + detail = f"{detail}; stderr_log={log_path}" + return ( + f"Distributed debug logits requested for {case.name} phase={phase}, " + f"but {detail}" + ) + + diff --git a/tests/e2e/models/timm_repvgg/e2e_plugins/runners/image_classification.py b/tests/e2e/models/timm_repvgg/e2e_plugins/runners/image_classification.py new file mode 100644 index 0000000000..b67d3faf89 --- /dev/null +++ b/tests/e2e/models/timm_repvgg/e2e_plugins/runners/image_classification.py @@ -0,0 +1,191 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Image-classification strategy runner.""" + +from __future__ import annotations + +import json +import logging +import os +import re +import subprocess +import tempfile +import time +from pathlib import Path + +from .. import save_full_stderr, _case_artifact_dir +from ..contracts import E2ECase, RunContext, StageOutput, StageSpec + +logger = logging.getLogger(__name__) +PROJECT_DIR = Path(__file__).resolve().parents[6] +_MPI_TAGGED_STDOUT_RE = re.compile( + r"^\[[^\]]+,(?P\d+)\]:\s?(?P.*)$") +_MPI_STREAM_TAG_RE = re.compile(r"\[[^\]]+,\d+\]<(?:stdout|stderr)>:\s?") + + +def _distributed_runtime_config(case: E2ECase) -> dict: + config = case.metadata.get("distributed_runtime", {}) + return config if isinstance(config, dict) and config.get("enabled") else {} + + +def _extract_rank_zero_stdout(stdout: str) -> str: + rank0_lines: list[str] = [] + saw_tagged = False + for line in (stdout or "").splitlines(): + match = _MPI_TAGGED_STDOUT_RE.match(line) + if match is None: + continue + saw_tagged = True + if int(match.group("rank")) == 0: + rank0_lines.append(match.group("text")) + if saw_tagged: + return "\n".join(rank0_lines).strip() + return (stdout or "").strip() + + +def _strip_mpi_stream_tags(text: str) -> str: + return _MPI_STREAM_TAG_RE.sub("", text or "") + + +def _ensure_distributed_runtime_env( + case: E2ECase, + ctx: RunContext, + env: dict[str, str], +) -> None: + if not _distributed_runtime_config(case) or env.get("TRTMC_NCCL_RENDEZVOUS"): + return + root = ( + Path(_case_artifact_dir(ctx.artifacts_dir, case.name)) + if ctx.artifacts_dir else Path(tempfile.gettempdir()) + ) + root.mkdir(parents=True, exist_ok=True) + path = root / f"{case.name}.nccl_rendezvous.bin" + try: + path.unlink() + except FileNotFoundError: + pass + env["TRTMC_NCCL_RENDEZVOUS"] = str(path) + + +def _wrap_distributed_command( + cmd: list[str], + case: E2ECase, + env: dict[str, str], +) -> list[str]: + config = _distributed_runtime_config(case) + if not config: + return cmd + launcher = str(config.get("launcher", "mpirun") or "mpirun") + world_size = int(config.get("world_size", config.get("tp_size", 2)) or 2) + launcher_args = config.get("launcher_args") + if isinstance(launcher_args, list): + prefix = [launcher] + [str(arg) for arg in launcher_args] + else: + prefix = [launcher, "--tag-output", "-np", str(world_size)] + + export_env = config.get("export_env", ["LD_LIBRARY_PATH", "CUDA_VISIBLE_DEVICES"]) + if isinstance(export_env, list) and Path(launcher).name == "mpirun": + export_names = [str(name) for name in export_env] + if "TRTMC_NCCL_RENDEZVOUS" in env and "TRTMC_NCCL_RENDEZVOUS" not in export_names: + export_names.append("TRTMC_NCCL_RENDEZVOUS") + for name in export_names: + if name in env: + prefix.extend(["-x", name]) + + return prefix + cmd + + +class ImageClassificationRunner: + @property + def strategy_name(self) -> str: + return "image_classification" + + def run_stage( + self, + case: E2ECase, + stage: StageSpec, + ctx: RunContext, + ) -> StageOutput: + bundle_path = os.path.join(ctx.engine_dir, case.bundle) + image_path = self._resolve_image_path(case, ctx) + cmd = [ + ctx.binary_path, + "classify", + bundle_path, + "--image", + image_path or "", + ] + runtime_cli_python = ctx.runtime_cli_hf_python() + if runtime_cli_python: + cmd.extend(["--hf-python", runtime_cli_python]) + if ctx.model_plugin_dir: + cmd.extend(["--model-plugin-dir", ctx.model_plugin_dir]) + + env = dict(os.environ) + if ctx.ld_library_path: + env["LD_LIBRARY_PATH"] = ctx.ld_library_path + _ensure_distributed_runtime_env(case, ctx, env) + cmd = _wrap_distributed_command(cmd, case, env) + + logger.info("Running image classification: %s", " ".join(cmd)) + t0 = time.monotonic() + result = subprocess.run( + cmd, capture_output=True, text=True, env=env, timeout=600) + elapsed = time.monotonic() - t0 + + if result.returncode != 0: + truncated, log_path = save_full_stderr( + result.stderr, ctx.artifacts_dir or "", + "image_classification", case.name) + msg = ( + f"Image classification failed (rc={result.returncode}): " + f"{truncated}" + ) + if log_path: + msg += f" (full stderr: {log_path})" + raise RuntimeError(msg) + + stdout = _extract_rank_zero_stdout(result.stdout) + try: + data = json.loads(stdout) + except json.JSONDecodeError: + data = {"raw_output": stdout} + + stderr_truncated, stderr_log = save_full_stderr( + result.stderr or "", ctx.artifacts_dir or "", + "image_classification", case.name) + metadata = { + "command": cmd, + "returncode": result.returncode, + "stdout": _strip_mpi_stream_tags(result.stdout or ""), + "stderr": _strip_mpi_stream_tags(stderr_truncated), + } + if stderr_log: + metadata["stderr_log"] = stderr_log + + return StageOutput( + stage_name=stage.name, + data=data, + timing_s=elapsed, + metadata=metadata, + ) + + def _resolve_image_path(self, case: E2ECase, ctx: RunContext) -> str | None: + image = ( + case.inputs.get("image") or case.inputs.get("test_image") + or case.inputs.get("image_path") + ) + if not image: + return None + path = Path(image) + if path.is_absolute(): + return str(path) + for base in (ctx.engine_dir, str(PROJECT_DIR), str(PROJECT_DIR / "tests" / "e2e")): + candidate = Path(base) / image + if candidate.is_file(): + return str(candidate) + return str(path) + + +plugin = ImageClassificationRunner() diff --git a/tests/e2e/models/timm_repvgg/e2e_plugins/runners/vl_debug_runner.py b/tests/e2e/models/timm_repvgg/e2e_plugins/runners/vl_debug_runner.py new file mode 100644 index 0000000000..deea0d8873 --- /dev/null +++ b/tests/e2e/models/timm_repvgg/e2e_plugins/runners/vl_debug_runner.py @@ -0,0 +1,1358 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Model-owned VL debug runner used by this E2E model plugin. + +This file intentionally duplicates the Python VL debug path so changes to one +model's runner do not couple sibling model E2E plugins. +""" + +from __future__ import annotations + +import ctypes +import os +import tempfile +import time +import warnings +from typing import Any + +import numpy as np +from tensorrt_model_connect import trt_compat + +trt = trt_compat.get_trt() if trt_compat.is_available() else None + +try: + from cuda.bindings import runtime as cudart +except ImportError: + try: + from cuda import cudart # type: ignore[no-redef] + except ImportError: # pragma: no cover - TRT-free test envs + cudart = None # type: ignore[assignment] + + +def _check_cuda(status): + if cudart is None: + raise RuntimeError("cuda-python is required for VL debug runner execution") + if hasattr(cudart, "cudaError_t"): + success = cudart.cudaError_t.cudaSuccess + else: + success = 0 + if status != success: + raise RuntimeError(f"CUDA error: {status}") + + +def _trt_nptype_safe(dtype: trt.DataType): + try: + return trt.nptype(dtype) + except TypeError: + if dtype == trt.bfloat16: + return np.uint16 + raise + + +def _trt_itemsize(dtype: trt.DataType) -> int: + return np.dtype(_trt_nptype_safe(dtype)).itemsize + + +def _require_trt_runtime() -> None: + if trt is None: + raise ImportError("tensorrt is required for VL debug runner execution") + if cudart is None: + raise ImportError("cuda-python is required for VL debug runner execution") + + +class TrtRunner: + """Device-resident TRT inference runner for debugging and diff testing. + + Keeps KV cache on-device. Only transfers token_id/position_id/mask + (H2D, ~1 KB) and logits + debug outputs (D2H) per step. Cache updates + are D2D memcpy. Matches the family-local decode cache behavior. + """ + + def __init__( + self, + engine_plan: bytes, + max_cache_length: int, + num_layers: int, + attention_size: int | None = None, + distributed_communicator: object | None = None, + ): + _require_trt_runtime() + self.max_cache_length = max_cache_length + self.num_layers = num_layers + self._distributed_communicator = distributed_communicator + + # Deserialize engine + logger = trt.Logger(trt.Logger.WARNING) + runtime = trt.Runtime(logger) + self.engine = runtime.deserialize_cuda_engine(engine_plan) + if self.engine is None: + raise RuntimeError("Failed to deserialize TRT engine") + self.context = self.engine.create_execution_context() + if distributed_communicator is not None: + set_communicator = getattr(self.context, "set_communicator", None) + if set_communicator is None: + raise RuntimeError( + "TensorRT distributed execution requires TRT 11.0+ " + "IExecutionContext.set_communicator" + ) + if not set_communicator(distributed_communicator): + raise RuntimeError("Failed to set TRT distributed communicator") + + # Dual-profile engines (built by build_dual_profile_decoder_engine) + # carry one prefill profile (profile 0, Sq dynamic) followed by one + # decode profile (profile 1, Sq=1). For per-step decode runs the + # debug_runner must select the decode profile and call + # set_input_shape on every dynamic input before each execute. We + # detect dual-profile via num_optimization_profiles > 1 and the + # presence of -1 dims on token_id / position_id / attention_mask. + self._dynamic_inputs: list[str] = [] + self._is_dual_profile = self.engine.num_optimization_profiles > 1 + for input_name in ("token_id", "position_id", "attention_mask"): + try: + shape = tuple(self.engine.get_tensor_shape(input_name)) + except Exception: + continue + if any(d < 0 for d in shape): + self._dynamic_inputs.append(input_name) + if self._is_dual_profile: + # Profile 1 = decode (Sq=1). step() always runs single-token, so + # we lock the context to that profile once and never switch. + err, decode_stream = cudart.cudaStreamCreate() + _check_cuda(err) + self.context.set_optimization_profile_async(1, decode_stream) + cudart.cudaStreamSynchronize(decode_stream) + cudart.cudaStreamDestroy(decode_stream) + + # Auto-detect attention_size from cache_k_0 shape + if attention_size is None: + cache_shape = tuple(self.engine.get_tensor_shape("cache_k_0")) + attention_size = cache_shape[1] # (max_cache_length, attention_size) + self.attention_size = attention_size + + # Detect cache element size from engine dtype (fp16=2, fp32=4) + cache_dtype = self.engine.get_tensor_dtype("cache_k_0") + self._cache_elem_bytes = _trt_itemsize(cache_dtype) + + # Create CUDA stream + err, self.stream = cudart.cudaStreamCreate() + _check_cuda(err) + + self.cache_length = 0 + attention_window = max_cache_length + 1 + row_bytes = self.attention_size * self._cache_elem_bytes + self._row_bytes = row_bytes + + # Discover IO tensor metadata and identify debug/extra outputs + self._output_names: list[str] = [] + self._output_shapes: dict[str, tuple] = {} + self._debug_output_names: list[str] = [] + for i in range(self.engine.num_io_tensors): + name = self.engine.get_tensor_name(i) + mode = self.engine.get_tensor_mode(name) + if mode == trt.TensorIOMode.OUTPUT: + shape = tuple(self.engine.get_tensor_shape(name)) + self._output_names.append(name) + self._output_shapes[name] = shape + # Debug outputs: anything that's not logits/present_k/present_v + if (name != "logits" + and not name.startswith("present_k_") + and not name.startswith("present_v_")): + self._debug_output_names.append(name) + + # --- Persistent device cache buffers (not copied per step) --- + cache_bytes = max_cache_length * row_bytes + self._d_cache_k: list[int] = [] + self._d_cache_v: list[int] = [] + for _ in range(num_layers): + err, dk = cudart.cudaMalloc(cache_bytes) + _check_cuda(err) + self._d_cache_k.append(dk) + err, dv = cudart.cudaMalloc(cache_bytes) + _check_cuda(err) + self._d_cache_v.append(dv) + + # --- Device buffers for present_k/v outputs (single-row each) --- + self._d_present_k: list[int] = [] + self._d_present_v: list[int] = [] + for _ in range(num_layers): + err, pk = cudart.cudaMalloc(row_bytes) + _check_cuda(err) + self._d_present_k.append(pk) + err, pv = cudart.cudaMalloc(row_bytes) + _check_cuda(err) + self._d_present_v.append(pv) + + # --- Small I/O: device + host buffers --- + self._h_token_id = np.zeros((1,), dtype=np.int32) + self._h_position_id = np.zeros((1,), dtype=np.int32) + err, self._d_token_id = cudart.cudaMalloc(4) + _check_cuda(err) + err, self._d_position_id = cudart.cudaMalloc(4) + _check_cuda(err) + + # attention_mask + self._h_mask = np.zeros((1, attention_window), dtype=np.float32) + err, self._d_mask = cudart.cudaMalloc(attention_window * 4) + _check_cuda(err) + + # logits + logits_shape = tuple(self.engine.get_tensor_shape("logits")) + self._logits_numel = int(np.prod(logits_shape)) + self._h_logits = np.zeros(logits_shape, dtype=np.float32) + err, self._d_logits = cudart.cudaMalloc(self._logits_numel * 4) + _check_cuda(err) + + # VL embed input support + self._has_embed_input = False + self._d_input_embed = 0 + self._d_use_input_embed = 0 + self._h_input_embed: np.ndarray | None = None + self._h_use_input_embed: np.ndarray | None = None + for i in range(self.engine.num_io_tensors): + name = self.engine.get_tensor_name(i) + if name == "input_embed": + self._has_embed_input = True + embed_shape = tuple(self.engine.get_tensor_shape(name)) + embed_bytes = int(np.prod(embed_shape)) * 4 + self._h_input_embed = np.zeros(embed_shape, dtype=np.float32) + err, self._d_input_embed = cudart.cudaMalloc(embed_bytes) + _check_cuda(err) + elif name == "use_input_embed": + self._h_use_input_embed = np.zeros((1,), dtype=np.float32) + err, self._d_use_input_embed = cudart.cudaMalloc(4) + _check_cuda(err) + + # DeepStack inputs (auto-detected from engine bindings) + self._deepstack_names: list[str] = [] + self._d_deepstack: dict[str, int] = {} + self._h_deepstack: dict[str, np.ndarray] = {} + self._d_deepstack_active = 0 + self._h_deepstack_active: np.ndarray | None = None + for i in range(self.engine.num_io_tensors): + name = self.engine.get_tensor_name(i) + if name.startswith("deepstack_embed_"): + self._deepstack_names.append(name) + shape = tuple(self.engine.get_tensor_shape(name)) + nbytes = int(np.prod(shape)) * 4 + self._h_deepstack[name] = np.zeros(shape, dtype=np.float32) + err, d_ptr = cudart.cudaMalloc(nbytes) + _check_cuda(err) + self._d_deepstack[name] = d_ptr + elif name == "deepstack_active": + self._h_deepstack_active = np.zeros((1,), dtype=np.float32) + err, self._d_deepstack_active = cudart.cudaMalloc(4) + _check_cuda(err) + + # Debug output device/host buffers + self._d_debug: dict[str, int] = {} + self._h_debug: dict[str, np.ndarray] = {} + for name in self._debug_output_names: + shape = self._output_shapes[name] + dtype_trt = self.engine.get_tensor_dtype(name) + dtype_np = _trt_nptype_safe(dtype_trt) + nbytes = int(np.prod(shape)) * np.dtype(dtype_np).itemsize + err, d_ptr = cudart.cudaMalloc(nbytes) + _check_cuda(err) + self._d_debug[name] = d_ptr + self._h_debug[name] = np.zeros(shape, dtype=dtype_np) + + # Zero-init device cache + for i in range(num_layers): + _check_cuda(cudart.cudaMemsetAsync( + self._d_cache_k[i], 0, cache_bytes, self.stream)[0]) + _check_cuda(cudart.cudaMemsetAsync( + self._d_cache_v[i], 0, cache_bytes, self.stream)[0]) + cudart.cudaStreamSynchronize(self.stream) + + @property + def has_embed_input(self) -> bool: + """True if the engine has input_embed and use_input_embed inputs.""" + return self._has_embed_input + + def step( + self, + token_id: int, + input_embed: np.ndarray | None = None, + use_input_embed: float = 0.0, + deepstack_embeds: list[np.ndarray] | None = None, + deepstack_active: float = 0.0, + ) -> dict[str, np.ndarray]: + """Run one decode step (manages position and cache internally). + + Args: + token_id: Input token ID. + input_embed: Optional pre-computed embedding [1, hidden] for VL prefill. + use_input_embed: 0.0 = use token_id lookup, 1.0 = use input_embed. + deepstack_embeds: Optional per-level DeepStack embeddings for VL prefill. + deepstack_active: 0.0 = inactive, 1.0 = inject DeepStack. + + Returns: + Dict with 'logits' and any debug outputs (e.g. 'debug_hidden_0'). + """ + H2D = cudart.cudaMemcpyKind.cudaMemcpyHostToDevice + D2H = cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost + D2D = cudart.cudaMemcpyKind.cudaMemcpyDeviceToDevice + stream = self.stream + attention_window = self.max_cache_length + 1 + + # Build attention mask (matches C++ build_attention_mask exactly) + position_id = min(self.cache_length, self.max_cache_length) + self._h_mask[:] = -1e9 + valid = min(self.cache_length, self.max_cache_length) + self._h_mask[0, :valid] = 0.0 + self._h_mask[0, -1] = 0.0 + + # Prepare small host buffers + self._h_token_id[0] = token_id + self._h_position_id[0] = position_id + + # H2D: small inputs only + cudart.cudaMemcpyAsync( + self._d_token_id, self._h_token_id.ctypes.data, + 4, H2D, stream) + cudart.cudaMemcpyAsync( + self._d_position_id, self._h_position_id.ctypes.data, + 4, H2D, stream) + cudart.cudaMemcpyAsync( + self._d_mask, self._h_mask.ctypes.data, + attention_window * 4, H2D, stream) + + # VL embed_input support + if self._has_embed_input: + if input_embed is not None and use_input_embed > 0.5: + self._h_input_embed[:] = input_embed.astype(np.float32) + self._h_use_input_embed[0] = use_input_embed + else: + self._h_input_embed[:] = 0.0 + self._h_use_input_embed[0] = 0.0 + cudart.cudaMemcpyAsync( + self._d_input_embed, self._h_input_embed.ctypes.data, + self._h_input_embed.nbytes, H2D, stream) + cudart.cudaMemcpyAsync( + self._d_use_input_embed, self._h_use_input_embed.ctypes.data, + 4, H2D, stream) + + # DeepStack H2D transfers + if self._deepstack_names: + for idx, ds_name in enumerate(self._deepstack_names): + if (deepstack_embeds is not None and idx < len(deepstack_embeds) + and deepstack_active > 0.5): + self._h_deepstack[ds_name][:] = deepstack_embeds[idx].astype(np.float32) + else: + self._h_deepstack[ds_name][:] = 0.0 + cudart.cudaMemcpyAsync( + self._d_deepstack[ds_name], + self._h_deepstack[ds_name].ctypes.data, + self._h_deepstack[ds_name].nbytes, H2D, stream) + if self._h_deepstack_active is not None: + self._h_deepstack_active[0] = deepstack_active + cudart.cudaMemcpyAsync( + self._d_deepstack_active, + self._h_deepstack_active.ctypes.data, + 4, H2D, stream) + + # Set tensor addresses + self.context.set_tensor_address("token_id", self._d_token_id) + self.context.set_tensor_address("position_id", self._d_position_id) + self.context.set_tensor_address("attention_mask", self._d_mask) + self.context.set_tensor_address("logits", self._d_logits) + + if self._has_embed_input: + self.context.set_tensor_address("input_embed", self._d_input_embed) + self.context.set_tensor_address( + "use_input_embed", self._d_use_input_embed) + + # DeepStack tensor binding (zeroed by default, set during VL prefill) + for ds_name in self._deepstack_names: + self.context.set_tensor_address(ds_name, self._d_deepstack[ds_name]) + if self._d_deepstack_active: + self.context.set_tensor_address( + "deepstack_active", self._d_deepstack_active) + + for i in range(self.num_layers): + self.context.set_tensor_address(f"cache_k_{i}", self._d_cache_k[i]) + self.context.set_tensor_address(f"cache_v_{i}", self._d_cache_v[i]) + self.context.set_tensor_address(f"present_k_{i}", self._d_present_k[i]) + self.context.set_tensor_address(f"present_v_{i}", self._d_present_v[i]) + + for name in self._debug_output_names: + self.context.set_tensor_address(name, self._d_debug[name]) + + # Dual-profile engines need explicit shapes for the dynamic + # inputs every step. step() is single-token decode, so all three + # shapes are fixed: Sq=1 and K = max_cache_length + 1. + if self._dynamic_inputs: + for name in self._dynamic_inputs: + if name == "attention_mask": + self.context.set_input_shape(name, (1, attention_window)) + else: + self.context.set_input_shape(name, (1,)) + + # Execute + self.context.execute_async_v3(stream) + + # D2D cache update + row_bytes = self.attention_size * self._cache_elem_bytes + for i in range(self.num_layers): + for cache_buf, present_buf in [ + (self._d_cache_k[i], self._d_present_k[i]), + (self._d_cache_v[i], self._d_present_v[i]), + ]: + if self.cache_length < self.max_cache_length: + offset = self.cache_length * row_bytes + cudart.cudaMemcpyAsync( + cache_buf + offset, present_buf, + row_bytes, D2D, stream) + else: + cudart.cudaMemcpyAsync( + cache_buf, cache_buf + row_bytes, + (self.max_cache_length - 1) * row_bytes, + D2D, stream) + offset = (self.max_cache_length - 1) * row_bytes + cudart.cudaMemcpyAsync( + cache_buf + offset, present_buf, + row_bytes, D2D, stream) + + # D2H: logits + debug outputs + cudart.cudaMemcpyAsync( + self._h_logits.ctypes.data, self._d_logits, + self._logits_numel * 4, D2H, stream) + for name in self._debug_output_names: + h_buf = self._h_debug[name] + cudart.cudaMemcpyAsync( + h_buf.ctypes.data, self._d_debug[name], + h_buf.nbytes, D2H, stream) + + cudart.cudaStreamSynchronize(stream) + self.cache_length = min(self.cache_length + 1, self.max_cache_length) + + # Collect results + results: dict[str, np.ndarray] = { + "logits": self._h_logits.copy(), + } + for name in self._debug_output_names: + results[name] = self._h_debug[name].copy() + return results + + def reset(self): + """Zero all device cache buffers and reset cache_length.""" + cache_bytes = self.max_cache_length * self.attention_size * self._cache_elem_bytes + for i in range(self.num_layers): + _check_cuda(cudart.cudaMemsetAsync( + self._d_cache_k[i], 0, cache_bytes, self.stream)[0]) + _check_cuda(cudart.cudaMemsetAsync( + self._d_cache_v[i], 0, cache_bytes, self.stream)[0]) + cudart.cudaStreamSynchronize(self.stream) + self.cache_length = 0 + + def generate( + self, + input_ids: list[int], + max_new_tokens: int, + ) -> list[dict[str, np.ndarray]]: + """Run autoregressive generation. + + Args: + input_ids: Prompt token IDs. + max_new_tokens: Number of tokens to generate after the prompt. + + Returns: + List of per-step result dicts. Each contains 'logits' and debug + outputs. The list includes both prefill steps (one per input token) + and generation steps. + """ + all_results = [] + + # Prefill: process input tokens one by one + for tid in input_ids: + result = self.step(tid) + all_results.append(result) + + # Generate: autoregressive decoding + for _ in range(max_new_tokens): + last_logits = all_results[-1]["logits"].flatten() + next_token = int(np.argmax(last_logits)) + result = self.step(next_token) + all_results.append(result) + + return all_results + + def __del__(self): + if cudart is None: + return + if not hasattr(self, "_d_token_id"): + return + bufs = [self._d_token_id, self._d_position_id, self._d_mask, + self._d_logits] + bufs.extend(self._d_cache_k) + bufs.extend(self._d_cache_v) + bufs.extend(self._d_present_k) + bufs.extend(self._d_present_v) + if self._d_input_embed: + bufs.append(self._d_input_embed) + if self._d_use_input_embed: + bufs.append(self._d_use_input_embed) + for d_ptr in self._d_deepstack.values(): + bufs.append(d_ptr) + if self._d_deepstack_active: + bufs.append(self._d_deepstack_active) + for d_ptr in self._d_debug.values(): + bufs.append(d_ptr) + for d_ptr in bufs: + cudart.cudaFree(d_ptr) + if hasattr(self, "stream"): + cudart.cudaStreamDestroy(self.stream) + if hasattr(self, "context"): + del self.context + if hasattr(self, "engine"): + del self.engine + + +_NCCL_UNIQUE_ID_BYTES = 128 +_NCCL_SUCCESS = 0 + + +class _NcclUniqueId(ctypes.Structure): + _fields_ = [("internal", ctypes.c_char * _NCCL_UNIQUE_ID_BYTES)] + + +def _env_int(names: tuple[str, ...], default: int | None = None) -> int | None: + for name in names: + raw = os.environ.get(name) + if raw is None or raw == "": + continue + try: + return int(raw) + except ValueError: + continue + return default + + +def _mpi_rank_info_from_env() -> tuple[int, int]: + rank = _env_int(("OMPI_COMM_WORLD_RANK", "PMI_RANK", "PMIX_RANK", "RANK"), 0) + world_size = _env_int( + ("OMPI_COMM_WORLD_SIZE", "PMI_SIZE", "PMIX_SIZE", "WORLD_SIZE"), 1) + return int(rank or 0), int(world_size or 1) + + +def _default_nccl_rendezvous_path() -> str: + path = os.environ.get("TRTMC_NCCL_RENDEZVOUS") + if path: + return path + job_id = ( + os.environ.get("OMPI_COMM_WORLD_JOBID") + or os.environ.get("PMIX_NAMESPACE") + or os.environ.get("SLURM_JOB_ID") + or f"pid{os.getppid()}" + ) + safe_job_id = "".join(ch if ch.isalnum() or ch in "._-" else "_" for ch in job_id) + return os.path.join(tempfile.gettempdir(), f"trtmc_nccl_{safe_job_id}.bin") + + +def _load_nccl_library() -> ctypes.CDLL: + errors: list[str] = [] + for name in ("libnccl.so.2", "libnccl.so"): + try: + lib = ctypes.CDLL(name) + break + except OSError as exc: + errors.append(f"{name}: {exc}") + else: + raise RuntimeError("Unable to load NCCL library: " + "; ".join(errors)) + + lib.ncclGetUniqueId.argtypes = [ctypes.POINTER(_NcclUniqueId)] + lib.ncclGetUniqueId.restype = ctypes.c_int + lib.ncclCommInitRank.argtypes = [ + ctypes.POINTER(ctypes.c_void_p), + ctypes.c_int, + _NcclUniqueId, + ctypes.c_int, + ] + lib.ncclCommInitRank.restype = ctypes.c_int + lib.ncclCommDestroy.argtypes = [ctypes.c_void_p] + lib.ncclCommDestroy.restype = ctypes.c_int + lib.ncclGetErrorString.argtypes = [ctypes.c_int] + lib.ncclGetErrorString.restype = ctypes.c_char_p + return lib + + +def _nccl_error_string(lib: ctypes.CDLL, status: int) -> str: + try: + msg = lib.ncclGetErrorString(status) + except Exception: + msg = None + if msg: + return msg.decode("utf-8", errors="replace") + return f"NCCL error {status}" + + +def _capsule_from_pointer(ptr: int): + pycapsule_new = ctypes.pythonapi.PyCapsule_New + pycapsule_new.restype = ctypes.py_object + pycapsule_new.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p] + return pycapsule_new(ctypes.c_void_p(ptr), None, None) + + +class TensorParallelNcclGroup: + """Small NCCL group helper for TensorRT distributed debug execution. + + The caller must destroy TRT execution contexts before closing this group. + """ + + def __init__( + self, + world_size: int | None = None, + rendezvous_path: str | None = None, + timeout_s: float = 60.0, + set_device: bool = True, + ): + _require_trt_runtime() + self.rank, detected_world_size = _mpi_rank_info_from_env() + self.world_size = int(world_size or detected_world_size) + if self.world_size <= 1: + raise RuntimeError("TensorParallelNcclGroup requires world_size > 1") + if detected_world_size != self.world_size: + raise RuntimeError( + f"MPI world size {detected_world_size} does not match " + f"requested tensor parallel size {self.world_size}" + ) + if self.rank < 0 or self.rank >= self.world_size: + raise RuntimeError( + f"MPI rank {self.rank} is outside world size {self.world_size}") + + if set_device: + status = cudart.cudaSetDevice(self.rank) + _check_cuda(status[0] if isinstance(status, tuple) else status) + + self.rendezvous_path = rendezvous_path or _default_nccl_rendezvous_path() + self._lib = _load_nccl_library() + self._comm = ctypes.c_void_p() + unique_id = self._exchange_unique_id(timeout_s=timeout_s) + self._check( + self._lib.ncclCommInitRank( + ctypes.byref(self._comm), + self.world_size, + unique_id, + self.rank, + ), + "ncclCommInitRank", + ) + if not self._comm.value: + raise RuntimeError("NCCL returned a null communicator") + self._communicator_capsule = _capsule_from_pointer(int(self._comm.value)) + self._closed = False + + @property + def communicator(self): + """PyCapsule wrapping the ncclComm_t pointer for TensorRT Python.""" + return self._communicator_capsule + + def _check(self, status: int, op: str) -> None: + if int(status) != _NCCL_SUCCESS: + raise RuntimeError(f"{op} failed: {_nccl_error_string(self._lib, status)}") + + def _exchange_unique_id(self, timeout_s: float) -> _NcclUniqueId: + path = self.rendezvous_path + if self.rank == 0: + unique_id = _NcclUniqueId() + self._check(self._lib.ncclGetUniqueId(ctypes.byref(unique_id)), "ncclGetUniqueId") + os.makedirs(os.path.dirname(path) or ".", exist_ok=True) + tmp_path = f"{path}.tmp.{os.getpid()}" + with open(tmp_path, "wb") as f: + f.write(ctypes.string_at(ctypes.byref(unique_id), _NCCL_UNIQUE_ID_BYTES)) + os.replace(tmp_path, path) + return unique_id + + deadline = time.monotonic() + timeout_s + while True: + try: + with open(path, "rb") as f: + data = f.read() + if len(data) == _NCCL_UNIQUE_ID_BYTES: + unique_id = _NcclUniqueId() + ctypes.memmove(ctypes.byref(unique_id), data, _NCCL_UNIQUE_ID_BYTES) + return unique_id + except FileNotFoundError: + pass + if time.monotonic() > deadline: + raise TimeoutError( + f"Timed out waiting for NCCL rendezvous file {path!r}") + time.sleep(0.05) + + def close(self) -> None: + if getattr(self, "_closed", True): + return + self._closed = True + if self._comm.value: + self._check(self._lib.ncclCommDestroy(self._comm), "ncclCommDestroy") + self._comm = ctypes.c_void_p() + + def __enter__(self) -> "TensorParallelNcclGroup": + return self + + def __exit__(self, exc_type, exc, tb) -> None: + self.close() + + def __del__(self): + try: + self.close() + except Exception: + pass + +def load_engine_from_bundle( + bundle_path: str, + section_name: str = "engine_plan", +) -> tuple[bytes, dict]: + """Load engine plan bytes and metadata from a .bundle artifact. + + Returns: + (engine_plan_bytes, header_dict) + """ + import json + import struct + + with open(bundle_path, "rb") as f: + magic = f.read(8) + if magic != b"BUNDLE\x01\x00": + raise ValueError(f"Not a valid .bundle artifact: {bundle_path}") + header_len = struct.unpack(" tuple[bytes | None, dict]: + """Load vision engine plan bytes from a .bundle artifact. + + Returns: + (vision_engine_plan_bytes_or_None, header_dict) + """ + import json + import struct + + with open(bundle_path, "rb") as f: + magic = f.read(8) + if magic != b"BUNDLE\x01\x00": + raise ValueError(f"Not a valid .bundle artifact: {bundle_path}") + header_len = struct.unpack(" dict[str, np.ndarray]: + """Run a single forward pass through the vision encoder. + + Args: + **inputs: Named input arrays (e.g. patch_embeds=...). + + Returns: + Dict of output name -> numpy array. + """ + # Set input values + for name, value in inputs.items(): + if name in self._host_buffers: + self._host_buffers[name][:] = value.astype( + self._host_buffers[name].dtype) + + # Copy inputs to device + for i in range(self.engine.num_io_tensors): + name = self.engine.get_tensor_name(i) + mode = self.engine.get_tensor_mode(name) + self.context.set_tensor_address(name, self._device_buffers[name]) + if mode == trt.TensorIOMode.INPUT: + h_buf = self._host_buffers[name] + cudart.cudaMemcpyAsync( + self._device_buffers[name], + h_buf.ctypes.data, + h_buf.nbytes, + cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, + self.stream, + ) + + self.context.execute_async_v3(self.stream) + + # Copy outputs + results: dict[str, np.ndarray] = {} + for name in self._output_names: + h_buf = self._host_buffers[name] + cudart.cudaMemcpyAsync( + h_buf.ctypes.data, + self._device_buffers[name], + h_buf.nbytes, + cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, + self.stream, + ) + + cudart.cudaStreamSynchronize(self.stream) + + for name in self._output_names: + results[name] = self._host_buffers[name].copy() + + return results + + def __del__(self): + if cudart is None: + return + for d_ptr in self._device_buffers.values(): + cudart.cudaFree(d_ptr) + if hasattr(self, "stream"): + cudart.cudaStreamDestroy(self.stream) + +def load_section_from_bundle(bundle_path: str, section_name: str) -> bytes | None: + """Load a named section's raw bytes from a .bundle artifact. + + Returns None if the section doesn't exist. + """ + import json + import struct + + with open(bundle_path, "rb") as f: + magic = f.read(8) + if magic != b"BUNDLE\x01\x00": + raise ValueError(f"Not a valid .bundle artifact: {bundle_path}") + header_len = struct.unpack(" dict: + """Load and parse config.json from a .bundle artifact.""" + import json + data = load_section_from_bundle(bundle_path, "config.json") + if data is None: + return {} + return json.loads(data.decode("utf-8")) + + +def load_preprocessor_config_from_bundle(bundle_path: str) -> dict: + """Load and parse preprocessor_config.json from a .bundle artifact.""" + import json + data = load_section_from_bundle(bundle_path, "preprocessor_config.json") + if data is None: + return {} + return json.loads(data.decode("utf-8")) + +def _resolve_pil_interpolation(mode: str): + """Map interpolation mode string to PIL constant.""" + from PIL import Image + _map = { + "bicubic": Image.BICUBIC, + "bilinear": Image.BILINEAR, + "nearest": Image.NEAREST, + } + return _map.get(mode, Image.BICUBIC) + + +def _preprocess_merge_group_chw( + image_path: str, + fixed_image_size: int = 448, + temporal_patch_size: int = 2, + image_mean: tuple[float, ...] = (0.48145466, 0.4578275, 0.40821073), + image_std: tuple[float, ...] = (0.26862954, 0.26130258, 0.27577711), + patch_size: int = 14, + merge_size: int = 2, + interpolation: str = "bicubic", +) -> np.ndarray: + """Merge-group preprocessing: [C*T, H, W] with patch permutation. + + The TRT vision engine's Conv2D produces patches in raster order of the + input image. To match HF's pipeline (where patches come out in merge-group + order after the processor's reshape+transpose), we spatially rearrange + the image at the 14x14 patch level: patches that should be in merge-group + positions are placed at the corresponding raster positions. + + Channel layout: [R_t0, R_t1, G_t0, G_t1, B_t0, B_t1] matching Conv3D + weight layout [out, C, T, kH, kW] reshaped to Conv2D [out, C*T, kH, kW]. + """ + from PIL import Image + + resample = _resolve_pil_interpolation(interpolation) + img = Image.open(image_path).convert("RGB") + img = img.resize((fixed_image_size, fixed_image_size), resample) + img_np = np.array(img, dtype=np.float32) / 255.0 + + # Normalize per channel + mean = np.array(image_mean, dtype=np.float32) + std = np.array(image_std, dtype=np.float32) + img_np = (img_np - mean) / std + + # HWC -> CHW + img_chw = img_np.transpose(2, 0, 1) # [C, H, W] + C = img_chw.shape[0] + T = temporal_patch_size + H = W = fixed_image_size + grid_h = H // patch_size + grid_w = W // patch_size + + # Extract 14x14 patches from the original image: [gH, gW, C, pH, pW] + orig_patches = img_chw.reshape( + C, grid_h, patch_size, grid_w, patch_size + ).transpose(1, 3, 0, 2, 4) # [gH, gW, C, pH, pW] + + # Build merge-group ordering: maps merge-group index -> (orig_h, orig_w) + # Merge groups iterate: (mh, mw, dh, dw) with orig = (mh*2+dh, mw*2+dw) + merge_h = grid_h // merge_size + merge_w = grid_w // merge_size + merge_idx = np.zeros((grid_h * grid_w, 2), dtype=np.int32) + idx = 0 + for mh in range(merge_h): + for mw in range(merge_w): + for dh in range(merge_size): + for dw in range(merge_size): + merge_idx[idx] = [mh * merge_size + dh, mw * merge_size + dw] + idx += 1 + + # Place merge-group-ordered patches at raster positions in pseudo-image. + # Patch at raster position i in the pseudo-image gets the content from + # the original position merge_idx[i], so Conv2D output patch i matches + # HF's i-th merge-group-ordered patch. + pseudo_patches = np.zeros_like(orig_patches) # [gH, gW, C, pH, pW] + for i in range(grid_h * grid_w): + ph = i // grid_w + pw = i % grid_w + oh, ow = merge_idx[i] + pseudo_patches[ph, pw] = orig_patches[oh, ow] + + # Reconstruct pseudo-image: [C, H, W] + pseudo_img = pseudo_patches.transpose(2, 0, 3, 1, 4).reshape(C, H, W) + + # Temporal duplication with [C, T] channel layout: + # [R_t0, R_t1, G_t0, G_t1, B_t0, B_t1] + pixel_values = np.repeat(pseudo_img, T, axis=0) # [C*T, H, W] + return pixel_values.astype(np.float32) + + +def _preprocess_simple_chw( + image_path: str, + fixed_image_size: int = 448, + image_mean: tuple[float, ...] = (0.48145466, 0.4578275, 0.40821073), + image_std: tuple[float, ...] = (0.26862954, 0.26130258, 0.27577711), + interpolation: str = "bicubic", + **_kwargs: Any, +) -> np.ndarray: + """Standard resize + normalize preprocessing: [C, H, W]. + + No patch permutation, no temporal duplication. Works for standard + ViT-based VL models. + """ + from PIL import Image + + resample = _resolve_pil_interpolation(interpolation) + img = Image.open(image_path).convert("RGB") + img = img.resize((fixed_image_size, fixed_image_size), resample) + img_np = np.array(img, dtype=np.float32) / 255.0 + + mean = np.array(image_mean, dtype=np.float32) + std = np.array(image_std, dtype=np.float32) + img_np = (img_np - mean) / std + + # HWC -> CHW + return img_np.transpose(2, 0, 1).astype(np.float32) + + +def _preprocess_patchify_chw( + image_path: str, + fixed_image_size: int = 448, + image_mean: tuple[float, ...] = (0.5, 0.5, 0.5), + image_std: tuple[float, ...] = (0.5, 0.5, 0.5), + patch_size: int = 14, + interpolation: str = "bicubic", + **_kwargs: Any, +) -> tuple[np.ndarray, np.ndarray]: + """Patchified CHW preprocessing: [N, C, pH, pW] plus [1, 2] grid.""" + from PIL import Image + + if patch_size <= 0 or fixed_image_size % patch_size != 0: + raise ValueError( + "fixed_image_size must be divisible by patch_size") + + resample = _resolve_pil_interpolation(interpolation) + img = Image.open(image_path).convert("RGB") + img = img.resize((fixed_image_size, fixed_image_size), resample) + img_np = np.array(img, dtype=np.float32) / 255.0 + + mean = np.array(image_mean, dtype=np.float32) + std = np.array(image_std, dtype=np.float32) + img_np = (img_np - mean) / std + + img_chw = img_np.transpose(2, 0, 1) + channels = img_chw.shape[0] + grid_h = fixed_image_size // patch_size + grid_w = fixed_image_size // patch_size + pixel_values = img_chw.reshape( + channels, grid_h, patch_size, grid_w, patch_size + ).transpose(1, 3, 0, 2, 4).reshape( + grid_h * grid_w, channels, patch_size, patch_size + ) + image_grid_hws = np.array([[grid_h, grid_w]], dtype=np.int32) + return pixel_values.astype(np.float32), image_grid_hws + + +def _preprocess_center_crop_chw( + image_path: str, + fixed_image_size: int = 448, + image_mean: tuple[float, ...] = (0.48145466, 0.4578275, 0.40821073), + image_std: tuple[float, ...] = (0.26862954, 0.26130258, 0.27577711), + interpolation: str = "bicubic", + **_kwargs: Any, +) -> np.ndarray: + """Center-crop to square, then resize + normalize: [C, H, W]. + + For traditional CLIP and DINOv2-based VL models that center-crop + before resize. + """ + from PIL import Image + + resample = _resolve_pil_interpolation(interpolation) + img = Image.open(image_path).convert("RGB") + + # Center-crop to square + w, h = img.size + crop_size = min(w, h) + left = (w - crop_size) // 2 + top = (h - crop_size) // 2 + img = img.crop((left, top, left + crop_size, top + crop_size)) + + img = img.resize((fixed_image_size, fixed_image_size), resample) + img_np = np.array(img, dtype=np.float32) / 255.0 + + mean = np.array(image_mean, dtype=np.float32) + std = np.array(image_std, dtype=np.float32) + img_np = (img_np - mean) / std + + return img_np.transpose(2, 0, 1).astype(np.float32) + + +def _preprocess_aspect_preserve_chw( + image_path: str, + fixed_image_size: int = 448, + image_mean: tuple[float, ...] = (0.48145466, 0.4578275, 0.40821073), + image_std: tuple[float, ...] = (0.26862954, 0.26130258, 0.27577711), + interpolation: str = "bicubic", + **_kwargs: Any, +) -> np.ndarray: + """Aspect-ratio-preserving resize + zero-pad to square: [C, H, W]. + + Fits image into fixed_image_size x fixed_image_size without distortion, + padding the remainder with zeros. + """ + from PIL import Image + + resample = _resolve_pil_interpolation(interpolation) + img = Image.open(image_path).convert("RGB") + + # Compute scaled dimensions fitting inside target square + w, h = img.size + scale = fixed_image_size / max(w, h) + new_w = max(1, int(w * scale)) + new_h = max(1, int(h * scale)) + + img = img.resize((new_w, new_h), resample) + + # Zero-pad to target square (top-left aligned) + padded = Image.new("RGB", (fixed_image_size, fixed_image_size), (0, 0, 0)) + padded.paste(img, (0, 0)) + + img_np = np.array(padded, dtype=np.float32) / 255.0 + + mean = np.array(image_mean, dtype=np.float32) + std = np.array(image_std, dtype=np.float32) + img_np = (img_np - mean) / std + + return img_np.transpose(2, 0, 1).astype(np.float32) + + +def _preprocess_pad_center_chw( + image_path: str, + fixed_image_size: int = 448, + image_mean: tuple[float, ...] = (0.48145466, 0.4578275, 0.40821073), + image_std: tuple[float, ...] = (0.26862954, 0.26130258, 0.27577711), + interpolation: str = "bicubic", + **_kwargs: Any, +) -> np.ndarray: + """Aspect-ratio-preserving resize + centered zero-pad: [C, H, W].""" + from PIL import Image + + resample = _resolve_pil_interpolation(interpolation) + img = Image.open(image_path).convert("RGB") + + w, h = img.size + scale = fixed_image_size / max(w, h) + new_w = max(1, int(w * scale)) + new_h = max(1, int(h * scale)) + + img = img.resize((new_w, new_h), resample) + + padded = Image.new("RGB", (fixed_image_size, fixed_image_size), (0, 0, 0)) + left = (fixed_image_size - new_w) // 2 + top = (fixed_image_size - new_h) // 2 + padded.paste(img, (left, top)) + + img_np = np.array(padded, dtype=np.float32) / 255.0 + + mean = np.array(image_mean, dtype=np.float32) + std = np.array(image_std, dtype=np.float32) + img_np = (img_np - mean) / std + + return img_np.transpose(2, 0, 1).astype(np.float32) + + +def preprocess_image_inputs_for_trt( + image_path: str, + preprocessor_type: str = "merge_group_chw", + **kwargs: Any, +) -> dict[str, np.ndarray]: + """Load and preprocess image inputs for a TRT vision engine. + + Returns named arrays keyed by TensorRT input name. Most models only need + pixel_values; patchified preprocessing also returns image_grid_hws. + """ + temporal = kwargs.get("temporal_patch_size", 1) + + if preprocessor_type == "patchify_chw": + pixel_values, image_grid_hws = _preprocess_patchify_chw( + image_path, **kwargs) + return { + "pixel_values": pixel_values, + "image_grid_hws": image_grid_hws, + } + + if preprocessor_type == "simple_chw": + result = _preprocess_simple_chw(image_path, **kwargs) + if temporal > 1 and result.shape[0] < temporal * 3: + result = np.tile(result, (temporal, 1, 1)) + return {"pixel_values": result} + if preprocessor_type == "center_crop_chw": + result = _preprocess_center_crop_chw(image_path, **kwargs) + if temporal > 1 and result.shape[0] < temporal * 3: + result = np.tile(result, (temporal, 1, 1)) + return {"pixel_values": result} + if preprocessor_type == "aspect_preserve_chw": + result = _preprocess_aspect_preserve_chw(image_path, **kwargs) + if temporal > 1 and result.shape[0] < temporal * 3: + result = np.tile(result, (temporal, 1, 1)) + return {"pixel_values": result} + if preprocessor_type == "pad_center_chw": + result = _preprocess_pad_center_chw(image_path, **kwargs) + if temporal > 1 and result.shape[0] < temporal * 3: + result = np.tile(result, (temporal, 1, 1)) + return {"pixel_values": result} + if preprocessor_type != "merge_group_chw": + warnings.warn( + f"Unknown preprocessor_type {preprocessor_type!r}, " + f"falling back to merge_group_chw", + stacklevel=2, + ) + return {"pixel_values": _preprocess_merge_group_chw(image_path, **kwargs)} + + +def preprocess_image_for_trt( + image_path: str, + preprocessor_type: str = "merge_group_chw", + **kwargs: Any, +) -> np.ndarray: + """Load and preprocess an image for the TRT vision engine. + + Compatibility wrapper returning only pixel_values. Use + preprocess_image_inputs_for_trt when the engine has auxiliary inputs. + """ + return preprocess_image_inputs_for_trt( + image_path, preprocessor_type=preprocessor_type, **kwargs)["pixel_values"] + +class VLTrtRunner: + """Full VL pipeline runner combining vision encoder + text decoder. + + Runs: preprocess image -> vision TRT -> build prompt -> text TRT decode. + Matches the C++ VLBackendFastPath pipeline exactly. + """ + + def __init__( + self, + bundle_path: str, + tokenizer=None, + ): + self.bundle_path = bundle_path + self.config = load_config_from_bundle(bundle_path) + self.preproc_config = load_preprocessor_config_from_bundle(bundle_path) + + # Load text decoder engine + engine_plan, header = load_engine_from_bundle(bundle_path) + self.text_runner = TrtRunner( + engine_plan=engine_plan, + max_cache_length=header["max_cache_length"], + num_layers=header["num_layers"], + ) + + # Load vision engine + vision_plan, _ = load_vision_engine_from_bundle(bundle_path) + self.vision_runner = VisionTrtRunner(vision_plan) if vision_plan else None + + # VL config from bundle + self.image_token_id = self.config.get("image_token_id", -1) + self.num_image_pad_tokens = self.config.get("num_image_pad_tokens", 256) + self.vl_prompt_template = self.config.get("vl_prompt_template", "") + self.image_token_str = self.config.get("image_token_str", "") + self.fixed_image_size = self.config.get("fixed_image_size", 448) + self.preprocessor_type = self.config.get( + "preprocessor_type", "merge_group_chw") + + # Preprocessor config + self.temporal_patch_size = self.preproc_config.get( + "temporal_patch_size", self.config.get("temporal_patch_size", 2)) + self.patch_size = self.preproc_config.get( + "patch_size", self.config.get("patch_size", 14)) + self.merge_size = self.preproc_config.get( + "merge_size", self.config.get("merge_size", 2)) + self.image_mean = tuple(self.preproc_config.get( + "image_mean", self.config.get( + "image_mean", [0.48145466, 0.4578275, 0.40821073]))) + self.image_std = tuple(self.preproc_config.get( + "image_std", self.config.get( + "image_std", [0.26862954, 0.26130258, 0.27577711]))) + self.interpolation = self.config.get("interpolation", "bicubic") + + self.tokenizer = tokenizer + + def encode_image(self, image_path: str) -> np.ndarray: + """Run the vision encoder on a single image. Returns [N, dim] features. + + Only single-image input is supported. Pass a single path string. + """ + if isinstance(image_path, (list, tuple)): + raise NotImplementedError( + "Multi-image input is not yet supported. " + "Pass a single image path string.") + if self.vision_runner is None: + raise RuntimeError("No vision engine in bundle") + + vision_inputs = preprocess_image_inputs_for_trt( + image_path, + preprocessor_type=self.preprocessor_type, + fixed_image_size=self.fixed_image_size, + temporal_patch_size=self.temporal_patch_size, + image_mean=self.image_mean, + image_std=self.image_std, + patch_size=self.patch_size, + merge_size=self.merge_size, + interpolation=self.interpolation, + ) + results = self.vision_runner.encode(**vision_inputs) + return results["image_features"] + + def format_prompt(self, user_prompt: str) -> str: + """Format the VL prompt with image pad tokens.""" + image_pads = self.image_token_str * self.num_image_pad_tokens + result = self.vl_prompt_template + result = result.replace("{image_pads}", image_pads) + result = result.replace("{prompt}", user_prompt) + return result + + def generate_vl( + self, + input_ids: list[int], + image_features: np.ndarray, + max_new_tokens: int, + ) -> list[int]: + """Run VL generation with pre-computed image features. + + Matches C++ VLBackendFastPath::generate_vl exactly: + - During prefill, image_token_id tokens are replaced with image features. + - During decode, normal autoregressive generation. + """ + feat_idx = 0 + output_ids = list(input_ids) + + # Prefill: all but last token + for tid in input_ids[:-1]: + embed = None + use_embed = 0.0 + if tid == self.image_token_id and feat_idx < len(image_features): + embed = image_features[feat_idx:feat_idx+1] # [1, dim] + use_embed = 1.0 + feat_idx += 1 + self.text_runner.step(tid, input_embed=embed, use_input_embed=use_embed) + + # Last prefill token + last_tid = input_ids[-1] + embed = None + use_embed = 0.0 + if last_tid == self.image_token_id and feat_idx < len(image_features): + embed = image_features[feat_idx:feat_idx+1] + use_embed = 1.0 + feat_idx += 1 + result = self.text_runner.step(last_tid, input_embed=embed, use_input_embed=use_embed) + + # Decode + for _ in range(max_new_tokens): + logits = result["logits"].flatten() + next_token = int(np.argmax(logits)) + output_ids.append(next_token) + eos_ids = self.config.get("eos_token_id", []) + if isinstance(eos_ids, int): + eos_ids = [eos_ids] + if next_token in eos_ids: + break + result = self.text_runner.step(next_token) + + return output_ids diff --git a/tests/e2e/models/timm_repvgg/e2e_plugins/runtime_config.py b/tests/e2e/models/timm_repvgg/e2e_plugins/runtime_config.py new file mode 100644 index 0000000000..f9e763100e --- /dev/null +++ b/tests/e2e/models/timm_repvgg/e2e_plugins/runtime_config.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Helpers for manifest-driven runtime config overrides.""" + +from __future__ import annotations + +from collections.abc import Iterator +from typing import Any + +from .contracts import E2ECase + + +def _runtime_config(case: E2ECase) -> dict[str, Any]: + config = case.metadata.get("runtime_config") + if isinstance(config, dict): + return config + config = case.inputs.get("runtime_config") + if isinstance(config, dict): + return config + return {} + + +def _flatten(prefix: str, value: Any) -> Iterator[tuple[str, Any]]: + if isinstance(value, dict): + for key, nested in value.items(): + name = f"{prefix}.{key}" if prefix else str(key) + yield from _flatten(name, nested) + elif prefix: + yield prefix, value + + +def _format_value(value: Any) -> str: + if isinstance(value, bool): + return "true" if value else "false" + return str(value) + + +def runtime_config_set_tokens(case: E2ECase) -> list[str]: + """Return CLI --set tokens from a manifest runtime_config mapping.""" + return [f"{name}={_format_value(value)}" for name, value in _flatten("", _runtime_config(case))] + + +def runtime_config_get(case: E2ECase, dotted_name: str, default: Any = None) -> Any: + value: Any = _runtime_config(case) + for part in dotted_name.split("."): + if not isinstance(value, dict) or part not in value: + return default + value = value[part] + return value diff --git a/tests/e2e/models/timm_repvgg/manifests/repvgg-a2-rvgg-in1k.json b/tests/e2e/models/timm_repvgg/manifests/repvgg-a2-rvgg-in1k.json new file mode 100644 index 0000000000..c70448da64 --- /dev/null +++ b/tests/e2e/models/timm_repvgg/manifests/repvgg-a2-rvgg-in1k.json @@ -0,0 +1,47 @@ +{ + "name": "repvgg-a2-rvgg-in1k", + "hf_id": "timm/repvgg_a2.rvgg_in1k", + "bundle": "repvgg-a2-rvgg-in1k.bundle", + "family": "timm_repvgg", + "runtime_strategy": "timm_repvgg_image_classification", + "task_strategy": "image_classification", + "max_cache_length": 1, + "precision": "fp16", + "trust_remote_code": false, + "testcases": [ + { + "name": "repvgg-a2-rvgg-in1k", + "trace_id": "IT-E2E-TIMM-REPVGG-01", + "reference_family": "image_classification", + "user_contract": "image_classification", + "test_type": "image_classification", + "prompt": "", + "max_new_tokens": 0, + "reference_precision": "fp32", + "test_image": "data/test_img.jpeg", + "preflight_requirements": [ + { + "kind": "binary_exists", + "args": {}, + "gating": true + }, + { + "kind": "asset_exists", + "args": { + "path": "data/test_img.jpeg" + }, + "gating": true + }, + { + "kind": "python_module_available", + "args": { + "module": "timm", + "phase": "reference" + }, + "gating": true + } + ], + "core": true + } + ] +} diff --git a/tests/e2e/models/timm_repvgg/runner.py b/tests/e2e/models/timm_repvgg/runner.py new file mode 100644 index 0000000000..2cba5838f3 --- /dev/null +++ b/tests/e2e/models/timm_repvgg/runner.py @@ -0,0 +1,186 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Model-owned E2E runner for the timm_repvgg family.""" + +from __future__ import annotations + +import os +import subprocess +import sys +from contextlib import contextmanager +from pathlib import Path + +from tests.e2e_harness.model_runner import ( + model_names_for_dir, + run_model_e2e as run_model_manifest_e2e, +) + +_MODEL_DIR = Path(__file__).resolve().parent +_PROJECT_DIR = _MODEL_DIR.parents[3] +_WAIVES_FILE = _MODEL_DIR / "waives.txt" + + +def _resolve_binary(config) -> str: + cli_val = config.getoption("--trtmc-binary", default=None) + if cli_val: + return str(Path(cli_val).absolute()) + default = _PROJECT_DIR / "build" / "trtmc" + return str(default) if default.is_file() else "" + + +def _resolve_hf_python(config) -> str: + cli_val = config.getoption("--hf-python", default=None) + if cli_val: + return str(Path(cli_val).absolute()) + venv = _PROJECT_DIR / ".venv" / "bin" / "python" + if venv.is_file(): + return str(venv) + return sys.executable + + +def _resolve_engine_dir(config) -> str: + cli_val = config.getoption("--engine-dir", default=None) + if cli_val: + d = Path(cli_val) + else: + d = Path("/mnt/storage/tensorrt-model-connect/engines") + d.mkdir(parents=True, exist_ok=True) + return str(d) + + +def _resolve_model_plugin_dir(config) -> str: + cli_val = config.getoption("--model-plugin-dir", default=None) + return str(Path(cli_val).absolute()) if cli_val else "" + + +def _resolve_artifacts_dir(config) -> str: + cli_val = config.getoption("--e2e-artifacts-dir", default=None) + if cli_val: + return str(Path(cli_val)) + return str(Path("/tmp/e2e_artifacts") / _MODEL_DIR.name) + + +@contextmanager +def _model_plugin_dir_env(path: str): + old_value = os.environ.get("TRTMC_MODEL_PLUGIN_DIR") + if path: + os.environ["TRTMC_MODEL_PLUGIN_DIR"] = path + try: + yield + finally: + if old_value is None: + os.environ.pop("TRTMC_MODEL_PLUGIN_DIR", None) + else: + os.environ["TRTMC_MODEL_PLUGIN_DIR"] = old_value + + +def _resolve_ld_library_path() -> str: + try: + result = subprocess.run( + [ + sys.executable, + "-c", + "import importlib.util; s=importlib.util.find_spec('tensorrt_libs'); " + "print(s.submodule_search_locations[0])", + ], + capture_output=True, + text=True, + timeout=10, + ) + trt_lib_dir = result.stdout.strip() + except Exception: + trt_lib_dir = "" + base = os.environ.get("LD_LIBRARY_PATH", "") + nccl_lib_dir = os.environ.get("TRTMC_NCCL_LIB_DIR", "") + parts = [p for p in [nccl_lib_dir, trt_lib_dir, "/usr/local/cuda/lib64", base] if p] + return ":".join(parts) + + +def _load_waives(platform: str = "") -> dict[str, tuple[str, str]]: + waives: dict[str, tuple[str, str]] = {} + if not _WAIVES_FILE.is_file(): + return waives + + platform = platform.strip() + with open(_WAIVES_FILE, encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line or line.startswith("#"): + continue + parts = line.split(None, 2) + if len(parts) < 2: + continue + + name_part = parts[0] + action = parts[1].upper() + reason = parts[2] if len(parts) > 2 else "" + if action not in ("SKIP", "XFAIL"): + continue + + if "/" in name_part: + plat, model_name = name_part.split("/", 1) + if plat != platform: + continue + else: + model_name = name_part + waives[model_name] = (action, reason) + return waives + + +def _is_multi_device_case(case) -> bool: + metadata = case.metadata or {} + return str(metadata.get("ci_tier", "") or "") == "multi_device" + + +def _parse_e2e_model_filters(values: list[str] | None) -> set[str]: + filters: set[str] = set() + for raw in values or []: + for item in str(raw).split(","): + item = item.strip() + if item: + filters.add(item) + return filters + + +def _case_matches_e2e_model(case, filters: set[str]) -> bool: + if not filters: + return True + metadata = case.metadata or {} + fields = { + case.name, + case.family, + case.runtime_strategy, + case.task_strategy, + Path(case.hf_id).name if case.hf_id else "", + str(metadata.get("family", "")), + str(metadata.get("runtime_strategy", "")), + } + return bool(filters & {field for field in fields if field}) + + +def model_case_names(config=None) -> list[str]: + return model_names_for_dir( + config=config, + model_dir=_MODEL_DIR, + case_matches_model=_case_matches_e2e_model, + is_multi_device_case=_is_multi_device_case, + ) + + +def run_model_e2e(case_name: str, request) -> None: + run_model_manifest_e2e( + model_name=case_name, + request=request, + model_dir=_MODEL_DIR, + load_waives=_load_waives, + case_matches_model=_case_matches_e2e_model, + is_multi_device_case=_is_multi_device_case, + resolve_hf_python=_resolve_hf_python, + resolve_artifacts_dir=_resolve_artifacts_dir, + resolve_binary=_resolve_binary, + resolve_ld_library_path=_resolve_ld_library_path, + resolve_engine_dir=_resolve_engine_dir, + resolve_model_plugin_dir=_resolve_model_plugin_dir, + model_plugin_dir_env=_model_plugin_dir_env, + ) diff --git a/tests/e2e/models/timm_repvgg/test_timm_repvgg_e2e.py b/tests/e2e/models/timm_repvgg/test_timm_repvgg_e2e.py new file mode 100644 index 0000000000..630554e055 --- /dev/null +++ b/tests/e2e/models/timm_repvgg/test_timm_repvgg_e2e.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Model-owned E2E entrypoint for the timm_repvgg family.""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path + +import pytest + +_RUNNER_PATH = Path(__file__).with_name("runner.py") +_SPEC = importlib.util.spec_from_file_location( + f"{Path(__file__).resolve().parent.name}_e2e_runner", + _RUNNER_PATH, +) +assert _SPEC is not None and _SPEC.loader is not None +_runner = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(_runner) + + +def pytest_generate_tests(metafunc): + if "case_name" in metafunc.fixturenames: + case_names = _runner.model_case_names(metafunc.config) + if not case_names: + pytest.skip("No model manifests selected", allow_module_level=True) + metafunc.parametrize("case_name", case_names) + + +def test_model_e2e(case_name: str, request) -> None: + _runner.run_model_e2e(case_name, request) diff --git a/tests/e2e/models/timm_repvgg/test_timm_repvgg_family_plugin.py b/tests/e2e/models/timm_repvgg/test_timm_repvgg_family_plugin.py new file mode 100644 index 0000000000..f55bbaace8 --- /dev/null +++ b/tests/e2e/models/timm_repvgg/test_timm_repvgg_family_plugin.py @@ -0,0 +1,183 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the timm RepVGG image-classification family plugin.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import pytest + +pytest.importorskip("tensorrt", reason="TensorRT is required for family builder tests") + + +try: + from safetensors.numpy import save_file + from tensorrt_model_connect.config import ModelConfig + from tensorrt_model_connect.families.timm_repvgg import plugin +except (ImportError, ModuleNotFoundError): + pytest.skip("tensorrt_model_connect requires TensorRT", allow_module_level=True) + + +def _rand(*shape: int) -> np.ndarray: + return np.random.RandomState(7).randn(*shape).astype(np.float32) + + +def _bn(t: dict, prefix: str, ch: int, *, unit: bool = False) -> None: + t[f"{prefix}.weight"] = np.ones(ch, dtype=np.float32) if unit else _rand(ch) + t[f"{prefix}.bias"] = np.zeros(ch, dtype=np.float32) if unit else _rand(ch) + t[f"{prefix}.running_mean"] = np.zeros(ch, dtype=np.float32) if unit else _rand(ch) + t[f"{prefix}.running_var"] = np.ones(ch, dtype=np.float32) if unit else np.abs(_rand(ch)) + 1.0 + + +def _branch(t: dict, prefix: str, out_ch: int, in_ch: int, k: int, *, unit: bool = False) -> None: + t[f"{prefix}.conv.weight"] = ( + np.zeros((out_ch, in_ch, k, k), dtype=np.float32) if unit else _rand(out_ch, in_ch, k, k) + ) + _bn(t, f"{prefix}.bn", out_ch, unit=unit) + + +def _write_tiny_repvgg(tmp_path: Path, *, blocks_per_stage: tuple[int, ...] = (2, 1)) -> None: + """Write a miniature RepVGG in its multi-branch training form.""" + classes = 5 + config = { + "architecture": "repvgg_a2", + "num_classes": classes, + "num_features": 16, + "pretrained_cfg": { + "input_size": [3, 224, 224], + "mean": [0.485, 0.456, 0.406], + "std": [0.229, 0.224, 0.225], + "crop_pct": 0.875, + "interpolation": "bilinear", + }, + } + (tmp_path / "config.json").write_text(json.dumps(config)) + + t: dict[str, np.ndarray] = {} + stem = 8 + _branch(t, "stem.conv_kxk", stem, 3, 3) + _branch(t, "stem.conv_1x1", stem, 3, 1) + + ch = stem + for stage, count in enumerate(blocks_per_stage): + out = 16 + for index in range(count): + p = f"stages.{stage}.{index}" + in_ch = ch if index == 0 else out + _branch(t, f"{p}.conv_kxk", out, in_ch, 3) + _branch(t, f"{p}.conv_1x1", out, in_ch, 1) + if index > 0: + # identity only exists where the shape is unchanged + _bn(t, f"{p}.identity", out) + ch = out + + t["head.fc.weight"] = _rand(classes, ch) + t["head.fc.bias"] = _rand(classes) + save_file(t, str(tmp_path / "model.safetensors")) + + +def test_model_config_uses_timm_architecture_when_model_type_absent(tmp_path: Path): + _write_tiny_repvgg(tmp_path) + cfg = ModelConfig.from_dir(tmp_path) + + assert cfg.model_type == "repvgg_a2" + assert plugin.matches(cfg.model_type) + + +@pytest.mark.parametrize("model_type", ["repvgg_a2", "repvgg_b0", "timm_repvgg"]) +def test_plugin_matches_repvgg_variants(model_type: str): + assert plugin.matches(model_type) + + +@pytest.mark.parametrize("model_type", ["resnet50", "vgg16", "ghostnet_100", ""]) +def test_plugin_rejects_unrelated_model_types(model_type: str): + assert not plugin.matches(model_type) + + +def test_stride_is_derived_from_the_identity_branch(tmp_path: Path): + """Only a block without an identity branch may change shape.""" + _write_tiny_repvgg(tmp_path, blocks_per_stage=(2, 1)) + cfg = ModelConfig.from_dir(tmp_path) + + plugin.load_weights(str(tmp_path), cfg) + blocks = cfg.raw["_timm_repvgg_config"]["blocks"] + + assert [b["has_identity"] for b in blocks] == [False, True, False] + assert [b["stride"] for b in blocks] == [2, 1, 2] + + +def test_load_weights_fuses_every_block_to_one_kernel(tmp_path: Path): + """Reparameterisation replaces three branches with a single 3x3 convolution.""" + _write_tiny_repvgg(tmp_path) + cfg = ModelConfig.from_dir(tmp_path) + + weights = plugin.load_weights(str(tmp_path), cfg) + + assert weights["stem.weight"].shape == (8, 3, 3, 3) + assert weights["stem.bias"].shape == (8,) + for block in cfg.raw["_timm_repvgg_config"]["blocks"]: + assert weights[f"{block['prefix']}.weight"].shape[-2:] == (3, 3) + # the multi-branch names must not survive into the built weights + assert not any("conv_1x1" in key or "identity" in key for key in weights) + + +def test_identity_branch_folds_to_a_centre_tap(tmp_path: Path): + """A unit batch norm on the identity path must contribute exactly 1 at the centre.""" + _write_tiny_repvgg(tmp_path, blocks_per_stage=(2,)) + from safetensors.numpy import load_file + + tensors = load_file(str(tmp_path / "model.safetensors")) + p = "stages.0.1" + # zero both convolutions and make the identity batch norm a pass-through + for leaf, k in ((f"{p}.conv_kxk", 3), (f"{p}.conv_1x1", 1)): + tensors[f"{leaf}.conv.weight"] = np.zeros_like(tensors[f"{leaf}.conv.weight"]) + ch = tensors[f"{leaf}.bn.weight"].shape[0] + tensors[f"{leaf}.bn.weight"] = np.zeros(ch, dtype=np.float32) + tensors[f"{leaf}.bn.bias"] = np.zeros(ch, dtype=np.float32) + tensors[f"{leaf}.bn.running_mean"] = np.zeros(ch, dtype=np.float32) + tensors[f"{leaf}.bn.running_var"] = np.ones(ch, dtype=np.float32) + ch = tensors[f"{p}.identity.weight"].shape[0] + tensors[f"{p}.identity.weight"] = np.ones(ch, dtype=np.float32) + tensors[f"{p}.identity.bias"] = np.zeros(ch, dtype=np.float32) + tensors[f"{p}.identity.running_mean"] = np.zeros(ch, dtype=np.float32) + tensors[f"{p}.identity.running_var"] = np.ones(ch, dtype=np.float32) + save_file(tensors, str(tmp_path / "model.safetensors")) + + cfg = ModelConfig.from_dir(tmp_path) + weights = plugin.load_weights(str(tmp_path), cfg) + fused = weights[f"{p}.weight"] + + expected = np.zeros_like(fused) + for channel in range(fused.shape[0]): + expected[channel, channel, 1, 1] = 1.0 + # 1/sqrt(1 + eps) is marginally below 1, so compare with a tolerance + np.testing.assert_allclose(fused, expected, atol=1e-4) + + +def test_load_weights_rejects_a_block_without_the_3x3_branch(tmp_path: Path): + _write_tiny_repvgg(tmp_path) + from safetensors.numpy import load_file + + kept = { + k: v + for k, v in load_file(str(tmp_path / "model.safetensors")).items() + if not k.startswith("stages.0.0.conv_kxk") + } + save_file(kept, str(tmp_path / "model.safetensors")) + cfg = ModelConfig.from_dir(tmp_path) + + with pytest.raises(ValueError, match="no conv_kxk"): + plugin.load_weights(str(tmp_path), cfg) + + +def test_build_engine_rejects_quantized_context(tmp_path: Path): + _write_tiny_repvgg(tmp_path) + cfg = ModelConfig.from_dir(tmp_path) + weights = plugin.load_weights(str(tmp_path), cfg) + + with pytest.raises(NotImplementedError, match="quantized"): + plugin.build_engine(cfg, weights, 0, quant_ctx=object()) diff --git a/tests/e2e/models/timm_repvgg/thresholds/repvgg-a2-rvgg-in1k.json b/tests/e2e/models/timm_repvgg/thresholds/repvgg-a2-rvgg-in1k.json new file mode 100644 index 0000000000..d4bbf65f5c --- /dev/null +++ b/tests/e2e/models/timm_repvgg/thresholds/repvgg-a2-rvgg-in1k.json @@ -0,0 +1,3 @@ +{ + "threshold_overrides": {} +} diff --git a/tests/runtime_strategy_matrix.yaml b/tests/runtime_strategy_matrix.yaml index 1ee3fad280..17dca24eac 100644 --- a/tests/runtime_strategy_matrix.yaml +++ b/tests/runtime_strategy_matrix.yaml @@ -61,6 +61,7 @@ "lerobot_act_action_chunk", "timm_vit_image_classification", "timm_resnet_image_classification", + "timm_repvgg_image_classification", "timm_vgg_image_classification", "whisper_speech_to_text", "canary_speech_to_text", @@ -946,6 +947,17 @@ "diff_framework_exemption": "No diff_framework check currently registers runtime_strategies=['timm_resnet_image_classification'].", "performance_mode": "single_pass" }, + "timm_repvgg_image_classification": { + "task_strategy": "image_classification", + "cli_commands": [ + "classify" + ], + "runner_class": "image_classification.ImageClassificationRunner", + "comparator_class": "image_classification.ImageClassificationComparator", + "diff_framework_check_classes": [], + "diff_framework_exemption": "No diff_framework check currently registers runtime_strategies=['timm_repvgg_image_classification'].", + "performance_mode": "single_pass" + }, "timm_vgg_image_classification": { "task_strategy": "image_classification", "cli_commands": [ diff --git a/tests/tools/test_model_plugin_encapsulation_static.py b/tests/tools/test_model_plugin_encapsulation_static.py index fe9111f74f..db06ad1b46 100644 --- a/tests/tools/test_model_plugin_encapsulation_static.py +++ b/tests/tools/test_model_plugin_encapsulation_static.py @@ -7875,7 +7875,7 @@ def test_hf_transformers_model_plugins_do_not_name_sibling_families() -> None: ), ), ( - {"timm_vit", "timm_resnet", "timm_vgg"}, + {"timm_vit", "timm_resnet", "timm_vgg", "timm_repvgg"}, ( "import timm", "timm.create_model", @@ -7987,7 +7987,7 @@ def test_single_family_e2e_task_sidecars_are_model_owned() -> None: """ cases = { "image_classification": { - "owners": {"timm_vit", "timm_resnet", "timm_vgg"}, + "owners": {"timm_vit", "timm_resnet", "timm_vgg", "timm_repvgg"}, "paths": ( "e2e_plugins/runners/image_classification.py", "e2e_plugins/comparators/image_classification.py", @@ -9197,14 +9197,14 @@ def test_generated_e2e_task_sidecars_are_task_owned() -> None: ( "runners", "image_classification.py", - {"timm_vit", "timm_resnet", "timm_vgg"}, + {"timm_vit", "timm_resnet", "timm_vgg", "timm_repvgg"}, "ImageClassificationRunner", "image_classification", ), ( "comparators", "image_classification.py", - {"timm_vit", "timm_resnet", "timm_vgg"}, + {"timm_vit", "timm_resnet", "timm_vgg", "timm_repvgg"}, "ImageClassificationComparator", "image_classification", ), diff --git a/tests/tools/test_perf_matrix.py b/tests/tools/test_perf_matrix.py index 3cd4742268..3186bd8f8f 100644 --- a/tests/tools/test_perf_matrix.py +++ b/tests/tools/test_perf_matrix.py @@ -76,6 +76,7 @@ def _suite_for_cases(cases, *, exclusions=None): "sana_wm.generate_image": "upstream-sana-wm", "segformer.segment": "hf-transformers-vision", "timesfm.solve": "pytorch-timeseries", + "timm_repvgg.classify": "hf-transformers-vision", "timm_resnet.classify": "hf-transformers-vision", "timm_vgg.classify": "hf-transformers-vision", "timm_vit.classify": "hf-transformers-vision", diff --git a/tests/validation/model_workloads.yaml b/tests/validation/model_workloads.yaml index 7cce3ec5ff..a9e952f130 100644 --- a/tests/validation/model_workloads.yaml +++ b/tests/validation/model_workloads.yaml @@ -262,6 +262,8 @@ models: workloads: [mmlu_five_shot_mcq] resnet50-a1-in1k: workloads: [imagenette_image_classification] + repvgg-a2-rvgg-in1k: + workloads: [imagenette_image_classification] riva-translate-4b: workloads: [flores200_en_fr_riva_translation_parity] roberta-base: diff --git a/tests/validation/workloads.yaml b/tests/validation/workloads.yaml index d39192a2fc..c78d510375 100644 --- a/tests/validation/workloads.yaml +++ b/tests/validation/workloads.yaml @@ -1259,10 +1259,12 @@ suites: runtime_strategies: - timm_vit_image_classification - timm_resnet_image_classification + - timm_repvgg_image_classification - timm_vgg_image_classification families: - timm_vit - timm_resnet + - timm_repvgg - timm_vgg user_contracts: - image_classification diff --git a/tools/legal_header_exceptions.toml b/tools/legal_header_exceptions.toml index 38d7f65509..0fe8cd795b 100644 --- a/tools/legal_header_exceptions.toml +++ b/tools/legal_header_exceptions.toml @@ -29,4 +29,4 @@ path = "tests/runtime_strategy_matrix.yaml" reason = "JSON-formatted test data is consumed by strict JSON parsers; comments would change parse behavior." license = "Apache-2.0" source = "https://github.com/NVIDIA/TensorRT-Model-Connect/blob/76c8164d97e645a996a210219ba635c8ec9a3453/tests/runtime_strategy_matrix.yaml" -sha256 = "5fa3e297e58ab4080044afda7f8950ed95c2bec263196f0848d51bb1d789403d" +sha256 = "603c0c513bf854b09191816a1b6dd10e3fb4ad98669d021128baa10289eb4025" diff --git a/website/data/hf-model-metadata.json b/website/data/hf-model-metadata.json index b000836fb3..47b46ed093 100644 --- a/website/data/hf-model-metadata.json +++ b/website/data/hf-model-metadata.json @@ -1311,6 +1311,17 @@ ], "architecture_source": "config.architectures" }, + { + "hf_id": "timm/repvgg_a2.rvgg_in1k", + "revision": "87d4d383cb45031cb9fa2fc8ddca73fd6649240f", + "revision_source": "resolved", + "metadata_file": "config.json", + "model_type": null, + "architectures": [ + "repvgg_a2" + ], + "architecture_source": "config.architecture" + }, { "hf_id": "timm/resnet50.a1_in1k", "revision": "767268603ca0cb0bfe326fa87277f19c419566ef", diff --git a/website/data/model-support-matrix.md b/website/data/model-support-matrix.md index 33711cbb0c..a91da3fb5c 100644 --- a/website/data/model-support-matrix.md +++ b/website/data/model-support-matrix.md @@ -103,6 +103,7 @@ SPDX-License-Identifier: Apache-2.0 | `bigcode/starcoder2-3b` | `starcoder2-3b` | `FP16` | None | — | 🟢 Green | | `google-t5/t5-small` | `t5-small` | `FP16` | None | — | 🟢 Green | | `google/timesfm-2.0-500m-pytorch` | `timesfm-2.0-500m-official` | `FP32` | None | — | 🟢 Green | +| `timm/repvgg_a2.rvgg_in1k` | `repvgg-a2-rvgg-in1k` | `FP16` | None | — | 🟢 Green | | `timm/resnet50.a1_in1k` | `resnet50-a1-in1k` | `FP16` | None | — | 🟢 Green | | `timm/vgg16.tv_in1k` | `vgg16-tv-in1k` | `FP16` | None | — | 🟢 Green | | `timm/vit_base_patch16_224.augreg_in21k_ft_in1k` | `timm-vit-base-p16-224-augreg-in21k-ft-in1k` | `FP16` | None | — | 🟢 Green | diff --git a/website/docs/features/runtime-strategies.md b/website/docs/features/runtime-strategies.md index e0c69ada4a..1ac45e9d37 100644 --- a/website/docs/features/runtime-strategies.md +++ b/website/docs/features/runtime-strategies.md @@ -23,7 +23,7 @@ selects that path before reading a native strategy. | Vision and multimodal | `qwen_vl_vision_language`, `internvl_vision_language`, `qwen3_omni_multimodal` | | Speech and audio | `whisper_speech_to_text`, `nemotron_speech_streaming_speech_to_text_rnnt`, `text_to_audio_bark`, `personaplex_speech_to_speech` | | Diffusion | `diffusion_flux`, `diffusion_wan`, `diffusion_wan2_2_ti2v`, `diffusion_qwen_image`, `diffusion_sana_wm` | -| Perception | `dinov3_image_feature_extraction`, `moge_monocular_geometry`, `segformer_segmentation`, `sam_prompted_segmentation`, `sam3_prompted_segmentation`, `timm_vit_image_classification`, `timm_resnet_image_classification`, `timm_vgg_image_classification` | +| Perception | `dinov3_image_feature_extraction`, `moge_monocular_geometry`, `segformer_segmentation`, `sam_prompted_segmentation`, `sam3_prompted_segmentation`, `timm_vit_image_classification`, `timm_resnet_image_classification`, `timm_repvgg_image_classification`, `timm_vgg_image_classification` | | Numeric operators | `chronos_bolt_trt`, `patchtsmixer_trt`, `patchtst_trt`, `timesfm_trt` | The complete live list is the union of the `runtime_strategies` arrays in the