diff --git a/benchmarks/performance/release.yaml b/benchmarks/performance/release.yaml index ce4d48307b..fbd75382f5 100644 --- a/benchmarks/performance/release.yaml +++ b/benchmarks/performance/release.yaml @@ -38,6 +38,10 @@ excluded_profiles: reason: >- The pinned Diffusers reference for MiniMax-H3 has not yet been integrated into the release performance runner. + - model: smollm3-3b + reason: >- + Dense SmolLM3 functional and reference-parity qualification is present, but + this change does not add a matching release-performance workload or receipt. entries: - id: albert.encode diff --git a/python/tensorrt_model_connect/families/smollm3/MODEL.toml b/python/tensorrt_model_connect/families/smollm3/MODEL.toml new file mode 100644 index 0000000000..c50e90a534 --- /dev/null +++ b/python/tensorrt_model_connect/families/smollm3/MODEL.toml @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +id = "smollm3" +plugin = "smollm3" +module = "plugin" +capabilities = ["split_decoder_roles", "debug_layer_outputs"] +aliases = [ + "smollm3", + "SmolLM3", +] +prefixes = [ + "smollm3", +] +architecture_patterns = [ + "SmolLM3ForCausalLM", +] +debug_runner = "debug_runner.py|runner_from_bundle" +default_build_route = "build_routing.py|prefer_native_default" +debug_runtime_strategies = [ + "smollm3_decoder_kv_cache", +] diff --git a/python/tensorrt_model_connect/families/smollm3/__init__.py b/python/tensorrt_model_connect/families/smollm3/__init__.py new file mode 100644 index 0000000000..2bcd0bb412 --- /dev/null +++ b/python/tensorrt_model_connect/families/smollm3/__init__.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import importlib +import sys +import types +from typing import Any + +_plugin = None + + +def _load_plugin_module(): + global _plugin + if _plugin is None: + _plugin = importlib.import_module(f"{__name__}.plugin") + globals().update({ + _name: _value + for _name, _value in vars(_plugin).items() + if not _name.startswith("__") + }) + return _plugin + + +def __getattr__(name: str) -> Any: + if name.startswith("__"): + raise AttributeError(name) + if name in {"graph_blocks", "graph_ops"}: + return importlib.import_module(f"{__name__}.{name}") + plugin_module = _load_plugin_module() + if name == "plugin": + return getattr(plugin_module, "plugin") + try: + return getattr(plugin_module, name) + except AttributeError: + raise AttributeError(name) from None + + +def __dir__() -> list[str]: + plugin_module = _load_plugin_module() + return sorted(set(globals()) | { + _name for _name in vars(plugin_module) if not _name.startswith("__") + }) + + +class _FamilyModule(types.ModuleType): + def __setattr__(self, name, value): + # Importlib publishes a directly imported plugin submodule on its parent. + # Keep the public package attribute bound to the FamilyPlugin instance. + if name == "plugin" and isinstance(value, types.ModuleType): + super().__setattr__("_plugin", value) + super().__setattr__("plugin", value.plugin) + return + super().__setattr__(name, value) + if ( + not name.startswith("__") + and name not in {"_plugin", "plugin"} + and not isinstance(value, types.ModuleType) + ): + setattr(_load_plugin_module(), name, value) + + +sys.modules[__name__].__class__ = _FamilyModule diff --git a/python/tensorrt_model_connect/families/smollm3/build_routing.py b/python/tensorrt_model_connect/families/smollm3/build_routing.py new file mode 100644 index 0000000000..20ae77d681 --- /dev/null +++ b/python/tensorrt_model_connect/families/smollm3/build_routing.py @@ -0,0 +1,379 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Routing contract for dense SmolLM3 models using TensorRT native KV cache.""" + +from __future__ import annotations + +import math +import operator + +_INT32_MAX = (1 << 31) - 1 +_UINT64_MAX = (1 << 64) - 1 + + +class NativeKvCapability: + """Small, loader-safe capability result (no dataclass dependency).""" + + __slots__ = ("applicable", "eligible", "reason") + + def __init__( + self, + applicable: bool, + eligible: bool, + reason: str, + ) -> None: + self.applicable = applicable + self.eligible = eligible + self.reason = reason + + +def _result( + *, + applicable: bool = True, + reasons: list[str] | tuple[str, ...] = (), +) -> NativeKvCapability: + return NativeKvCapability( + applicable, + applicable and not reasons, + "; ".join(reasons) or "supported", + ) + + +def _raw(config: object) -> dict: + value = getattr(config, "raw", {}) + return value if isinstance(value, dict) else {} + + +def _integer(value: object, name: str) -> int: + if isinstance(value, bool): + raise ValueError(f"{name} must be an integer") + try: + return int(operator.index(value)) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError(f"{name} must be an integer") from exc + + +def _positive(config: object, name: str) -> int: + value = _integer(getattr(config, name, None), name) + if value <= 0: + raise ValueError(f"{name} must be positive") + if value > _INT32_MAX: + raise ValueError(f"{name} exceeds TensorRT's int32 dimension limit") + return value + + +def resolved_head_dim(config: object) -> int: + """Return the explicit HF head width, or derive it when absent.""" + + raw = _raw(config) + explicit = raw.get("head_dim", getattr(config, "_head_dim", 0)) + if "head_dim" in raw or explicit not in (None, 0): + head_dim = _integer(explicit, "head_dim") + else: + hidden = _positive(config, "hidden_size") + heads = _positive(config, "num_attention_heads") + if hidden % heads: + raise ValueError( + "hidden_size must be divisible by num_attention_heads when " + "head_dim is absent" + ) + head_dim = hidden // heads + if not 0 < head_dim <= _INT32_MAX: + raise ValueError("head_dim must be a positive TensorRT dimension") + return head_dim + + +def _checked_product(label: str, *values: int) -> int: + product = 1 + for value in values: + if value <= 0 or product > _UINT64_MAX // value: + raise ValueError(f"native SmolLM3 KV {label} exceeds uint64") + product *= value + return product + + +def native_kv_cache_geometry( + config: object, + capacity: int, + *, + element_bytes: int = 2, +) -> tuple[int, int]: + """Return runtime byte geometry for the required full-context cache.""" + + capacity = _integer(capacity, "max_cache_length") + context = _positive(config, "max_position_embeddings") + if capacity != context: + raise ValueError( + "native SmolLM3 KV requires max_cache_length == " + f"max_position_embeddings ({context}), got {capacity}" + ) + row_bytes = _checked_product( + "row size", + 2, + _positive(config, "num_hidden_layers"), + _positive(config, "num_key_value_heads"), + resolved_head_dim(config), + _integer(element_bytes, "element_bytes"), + ) + return row_bytes, _checked_product("cache size", capacity, row_bytes) + + +def _enabled(value: object) -> bool: + return value not in (None, False, 0, "", (), [], {}) + + +def _validate_rope(raw: dict, reasons: list[str]) -> None: + parameters = raw.get("rope_parameters") + scaling = raw.get("rope_scaling") + if parameters is not None and scaling is not None: + reasons.append("RoPE configuration is ambiguous") + return + rope = parameters if parameters is not None else scaling + if rope is None: + return + if not isinstance(rope, dict): + reasons.append("RoPE configuration must be an object") + return + rope_type = str(rope.get("rope_type", rope.get("type", "default"))).lower() + if rope_type in ("", "default"): + if any( + key in rope + for key in ( + "factor", + "low_freq_factor", + "high_freq_factor", + "original_max_position_embeddings", + ) + ): + reasons.append("default RoPE must not contain scaling parameters") + return + if rope_type != "llama3": + # The family still builds these: they route to the standard decoder, + # which applies the scaling through an indexed cos/sin table. Only the + # native KV graph is limited to unscaled and llama3 RoPE. + reasons.append( + f"native SmolLM3 KV does not support rope_type={rope_type!r}") + return + required = ( + "factor", + "low_freq_factor", + "high_freq_factor", + "original_max_position_embeddings", + ) + if any(name not in rope for name in required): + reasons.append("llama3 RoPE is missing required scaling parameters") + return + try: + factor = float(rope["factor"]) + low = float(rope["low_freq_factor"]) + high = float(rope["high_freq_factor"]) + original = _integer( + rope["original_max_position_embeddings"], + "original_max_position_embeddings", + ) + except (TypeError, ValueError, OverflowError): + reasons.append("llama3 RoPE scaling parameters must be numeric") + return + if ( + not all(math.isfinite(value) for value in (factor, low, high)) + or factor < 1.0 + or low <= 0.0 + or high <= low + or original <= 0 + ): + reasons.append("llama3 RoPE scaling parameters are invalid") + + +def _validate_rope_layer_schedule( + config: object, + raw: dict, + reasons: list[str], +) -> None: + """Reject a malformed NoPE schedule at routing time. + + SmolLM3 marks NoPE layers through ``no_rope_layers`` (1 = the layer applies + RoPE, 0 = NoPE) or, when that list is absent, ``no_rope_layer_interval``. + Resolving it here keeps a bad schedule from surfacing as an exception deep + in the graph builder. + """ + if raw.get("no_rope_layers") is None and raw.get( + "no_rope_layer_interval" + ) is None: + return + try: + # Absolute, not relative: MODEL.toml points default_build_route at + # this file, and families/__init__.py loads it by path under a + # synthetic top-level name, so it has no package context. The + # file-loaded debug_runner modules import the package the same way. + from tensorrt_model_connect.families.smollm3.config import ( + resolve_rope_layer_schedule, + ) + + schedule = resolve_rope_layer_schedule(config) + except ValueError as exc: + reasons.append(str(exc)) + return + num_layers = int(getattr(config, "num_hidden_layers", 0) or 0) + if len(schedule) != num_layers: + reasons.append( + "NoPE schedule must cover every layer: got " + f"{len(schedule)} entries for {num_layers} layers" + ) + + +def native_kv_architecture_capability( + config: object, +) -> NativeKvCapability: + """Accept any model size that retains the dense SmolLM3 graph contract.""" + + model_type = str(getattr(config, "model_type", "")).lower() + if not model_type.startswith("smollm3"): + return _result(applicable=False) + if model_type != "smollm3": + return _result(reasons=["model_type must be exactly 'smollm3'"]) + + raw = _raw(config) + reasons: list[str] = [] + if tuple(getattr(config, "architectures", ()) or ()) != ( + "SmolLM3ForCausalLM", + ): + reasons.append("architectures must contain exactly SmolLM3ForCausalLM") + + try: + dimensions = { + name: _positive(config, name) + for name in ( + "vocab_size", + "hidden_size", + "intermediate_size", + "num_hidden_layers", + "num_attention_heads", + "num_key_value_heads", + "max_position_embeddings", + ) + } + head_dim = resolved_head_dim(config) + if dimensions["num_attention_heads"] % dimensions[ + "num_key_value_heads" + ]: + reasons.append( + "num_attention_heads must be divisible by " + "num_key_value_heads" + ) + if head_dim != 128: + reasons.append("native SmolLM3 attention requires head_dim=128") + except ValueError as exc: + reasons.append(str(exc)) + + if str(getattr(config, "hidden_act", "")).lower() != "silu": + reasons.append("native SmolLM3 requires hidden_act='silu'") + for name in ("rms_norm_eps", "rope_theta"): + try: + value = float(getattr(config, name)) + except (TypeError, ValueError, OverflowError): + value = 0.0 + if not math.isfinite(value) or value <= 0: + reasons.append(f"{name} must be finite and positive") + + unsupported_flags = ( + "attention_bias", + "mlp_bias", + "is_encoder_decoder", + "use_sliding_window", + "sliding_window", + "rope_interleaved", + "interleaved_rope", + "num_experts", + "num_local_experts", + "num_experts_per_tok", + ) + enabled = [name for name in unsupported_flags if _enabled(raw.get(name))] + if enabled: + reasons.append("unsupported SmolLM3 fields: " + ", ".join(enabled)) + # pretraining_tp is deliberately not gated. SmolLM3ForCausalLM neither + # defines nor reads it -- the field is absent from both the upstream + # configuration class and the modeling code -- so it describes nothing + # about the graph this family builds. The published checkpoint still + # carries pretraining_tp=2 from the Llama-derived template it started + # from, and rejecting that would route the only checkpoint this family + # targets away from the path its manifest declares. Llama gates it + # because Llama's own implementation reads it. + try: + if float(raw.get("partial_rotary_factor", 1.0)) != 1.0: + reasons.append("native SmolLM3 requires full rotary embeddings") + except (TypeError, ValueError, OverflowError): + reasons.append("partial_rotary_factor must be numeric") + layer_types = raw.get("layer_types") + if layer_types is not None and ( + not isinstance(layer_types, (list, tuple)) + or any(str(value).lower() != "full_attention" for value in layer_types) + ): + reasons.append("native SmolLM3 does not support hybrid layer types") + _validate_rope(raw, reasons) + _validate_rope_layer_schedule(config, raw, reasons) + return _result(reasons=reasons) + + +def native_kv_build_capability( + config: object, + *, + precision: str = "bf16", + max_cache_length: int | None = None, + parallel_enabled: bool | None = None, + dynamic_kv_cache: bool | None = None, + quantized: bool | None = None, + debug_layer_outputs: bool = False, +) -> NativeKvCapability: + """Apply deployment constraints once, after architecture routing.""" + + architecture = native_kv_architecture_capability(config) + if not architecture.eligible: + return architecture + + raw = _raw(config) + reasons: list[str] = [] + if str(precision).lower() != "bf16": + reasons.append("native SmolLM3 requires BF16") + if str(raw.get("_decoder_engine_layout", "split")) != "split": + reasons.append("native SmolLM3 requires split prefill/decode engines") + if raw.get("_rtx_build_requested"): + reasons.append("native SmolLM3 requires the standard TensorRT backend") + if parallel_enabled or raw.get("_parallel_build_enabled"): + reasons.append("native SmolLM3 does not support tensor parallel builds") + if ( + dynamic_kv_cache + or raw.get("_runtime_dynamic_kv_requested") + or raw.get("dynamic_kv_cache") + ): + reasons.append("native SmolLM3 uses one fixed physical KV capacity") + if ( + quantized + or raw.get("quantization_config") + or raw.get("_quantized_build_requested") + ): + reasons.append("native SmolLM3 does not support quantized builds") + if raw.get("_fp32_layers"): + reasons.append("native SmolLM3 does not support FP32 layer overrides") + if debug_layer_outputs: + reasons.append("native SmolLM3 does not support debug layer outputs") + try: + native_kv_cache_geometry( + config, + ( + int(getattr(config, "max_position_embeddings")) + if max_cache_length is None + else max_cache_length + ), + ) + except ValueError as exc: + reasons.append(str(exc)) + return _result(reasons=reasons) + + +def prefer_native_default( + config: object, +) -> bool: + """Route dense SmolLM3 to native KV without a user-facing build flag.""" + + return native_kv_architecture_capability(config).eligible diff --git a/python/tensorrt_model_connect/families/smollm3/checkpoint_mapper.py b/python/tensorrt_model_connect/families/smollm3/checkpoint_mapper.py new file mode 100644 index 0000000000..e668d1edb3 --- /dev/null +++ b/python/tensorrt_model_connect/families/smollm3/checkpoint_mapper.py @@ -0,0 +1,386 @@ +# 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 + +import os +from concurrent.futures import ThreadPoolExecutor, as_completed +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 + +from .config import ModelConfig + + +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 _layer_key(layer_idx: int, suffix: str, model_prefix: str = "model") -> str: + return f"{model_prefix}.layers.{layer_idx}.{suffix}" + + +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)) + + +def _repeat_head_norm(norm: np.ndarray, num_heads: int) -> np.ndarray: + """Repeat per-head norm [head_dim] -> [num_heads * head_dim].""" + return np.tile(norm, num_heads).astype(np.float32) + + +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] + """ + + +def load_standard_weights( + model_dir: str | Path, + config: ModelConfig, + *, + precision: str = "fp32", + fp32_layers: tuple[int, ...] = (), + model_prefix: str = "model", + embedding_key: str | None = None, + final_norm_key: str | None = None, + lm_head_key: str = "lm_head.weight", +) -> WeightDict: + """Load HF safetensors and map to standard weight dict.""" + model_dir = Path(model_dir) + readers = _open_safetensors(model_dir) + + hidden = config.hidden_size + vocab = config.vocab_size + num_layers = config.num_hidden_layers + num_heads = config.num_attention_heads + num_kv_heads = config.num_key_value_heads + target_dtype = _target_np_dtype(precision) + selected_fp32_layers = frozenset(int(layer) for layer in fp32_layers) + invalid_fp32_layers = sorted( + layer for layer in selected_fp32_layers + if layer < 0 or layer >= num_layers) + if invalid_fp32_layers: + raise ValueError( + f"fp32_layers contains out-of-range indices: {invalid_fp32_layers}") + if precision == "fp32": + selected_fp32_layers = frozenset() + + weights = WeightDict() + + # Embedding + if embedding_key is None: + embedding_key = f"{model_prefix}.embed_tokens.weight" + embedding = _load_tensor(readers, embedding_key) + assert embedding.shape == (vocab, hidden), ( + f"Embedding shape {embedding.shape} != ({vocab}, {hidden})") + weights["embedding"] = embedding.astype(target_dtype) + + def _load_layer(layer_idx: int) -> tuple[int, WeightDict, int, int]: + prefix = f"layer.{layer_idx}" + layer = WeightDict() + layer_precision = ( + "fp32" if layer_idx in selected_fp32_layers else precision) + layer_target_dtype = _target_np_dtype(layer_precision) + + # Norms + input_norm = _load_tensor( + readers, _layer_key(layer_idx, "input_layernorm.weight", model_prefix)) + post_norm = _load_tensor( + readers, + _layer_key(layer_idx, "post_attention_layernorm.weight", model_prefix)) + layer[f"{prefix}.input_norm"] = input_norm.astype(np.float32) + layer[f"{prefix}.post_attn_norm"] = post_norm.astype(np.float32) + + # Q/K/V/O projections + q_raw = _load_tensor( + readers, _layer_key(layer_idx, "self_attn.q_proj.weight", model_prefix)) + k_raw = _load_tensor( + readers, _layer_key(layer_idx, "self_attn.k_proj.weight", model_prefix)) + v_raw = _load_tensor( + readers, _layer_key(layer_idx, "self_attn.v_proj.weight", model_prefix)) + o_raw = _load_tensor( + readers, _layer_key(layer_idx, "self_attn.o_proj.weight", model_prefix)) + + q_hidden = q_raw.shape[0] + gate_raw = _load_tensor( + readers, _layer_key(layer_idx, "mlp.gate_proj.weight", model_prefix)) + layer_mlp_size = gate_raw.shape[0] + + # Transpose all projections [out, in] -> [in, out] + q_t = _transpose_2d(q_raw, "q_proj", precision=layer_precision) + k_t = _transpose_2d(k_raw, "k_proj", precision=layer_precision) + v_t = _transpose_2d(v_raw, "v_proj", precision=layer_precision) + o_t = _transpose_2d(o_raw, "o_proj", precision=layer_precision) + + layer[f"{prefix}.w_q"] = q_t + layer[f"{prefix}.w_k"] = k_t + layer[f"{prefix}.w_v"] = v_t + layer[f"{prefix}.w_o"] = o_t + + # Optional QKV biases (Qwen2 style) + q_bias_key = _layer_key(layer_idx, "self_attn.q_proj.bias", model_prefix) + k_bias_key = _layer_key(layer_idx, "self_attn.k_proj.bias", model_prefix) + v_bias_key = _layer_key(layer_idx, "self_attn.v_proj.bias", model_prefix) + if _has_tensor(readers, q_bias_key): + layer[f"{prefix}.q_bias"] = _load_tensor( + readers, q_bias_key).astype(layer_target_dtype) + if _has_tensor(readers, k_bias_key): + layer[f"{prefix}.k_bias"] = _load_tensor( + readers, k_bias_key).astype(layer_target_dtype) + if _has_tensor(readers, v_bias_key): + layer[f"{prefix}.v_bias"] = _load_tensor( + readers, v_bias_key).astype(layer_target_dtype) + + # Optional per-head q/k norm (Qwen3 style) + q_norm_key = _layer_key(layer_idx, "self_attn.q_norm.weight", model_prefix) + k_norm_key = _layer_key(layer_idx, "self_attn.k_norm.weight", model_prefix) + if _has_tensor(readers, q_norm_key): + layer[f"{prefix}.q_norm"] = _repeat_head_norm( + _load_tensor(readers, q_norm_key).astype(np.float32), + num_heads) + if _has_tensor(readers, k_norm_key): + layer[f"{prefix}.k_norm"] = _repeat_head_norm( + _load_tensor(readers, k_norm_key).astype(np.float32), + num_kv_heads) + + # MLP projections + up_raw = _load_tensor( + readers, _layer_key(layer_idx, "mlp.up_proj.weight", model_prefix)) + down_raw = _load_tensor( + readers, _layer_key(layer_idx, "mlp.down_proj.weight", model_prefix)) + + layer[f"{prefix}.w_gate"] = _transpose_2d( + gate_raw, "gate_proj", precision=layer_precision) + layer[f"{prefix}.w_up"] = _transpose_2d( + up_raw, "up_proj", precision=layer_precision) + layer[f"{prefix}.w_down"] = _transpose_2d( + down_raw, "down_proj", precision=layer_precision) + + return layer_idx, layer, q_hidden, layer_mlp_size + + layer_results: list[tuple[int, WeightDict, int, int] | None] = [None] * num_layers + max_workers = min(8, max(1, os.cpu_count() or 1)) + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = [executor.submit(_load_layer, i) for i in range(num_layers)] + for future in as_completed(futures): + layer_idx, layer, attention_size, mlp_size = future.result() + layer_results[layer_idx] = (layer_idx, layer, attention_size, mlp_size) + + attention_size = 0 + kv_attention_size = 0 + mlp_size = 0 + for result in layer_results: + if result is None: + continue + _layer_idx, layer, layer_attention_size, layer_mlp_size = result + weights.update(layer) + if attention_size == 0: + attention_size = layer_attention_size + first_k = layer[f"layer.{_layer_idx}.w_k"] + kv_attention_size = int(first_k.shape[1]) + if mlp_size == 0: + mlp_size = layer_mlp_size + + # Final norm + if final_norm_key is None: + final_norm_key = f"{model_prefix}.norm.weight" + if _has_tensor(readers, final_norm_key): + weights["final_norm"] = _load_tensor( + readers, final_norm_key).astype(np.float32) + else: + weights["final_norm"] = np.ones(hidden, dtype=np.float32) + + # LM head + if _has_tensor(readers, lm_head_key): + weights["w_out"] = _transpose_2d( + _load_tensor(readers, lm_head_key), "lm_head", precision=precision) + else: + # Tied embeddings + weights["w_out"] = _transpose_2d(embedding.copy(), "embedding_tied", + precision=precision) + + weights["_attention_size"] = attention_size # type: ignore[assignment] + weights["_kv_attention_size"] = kv_attention_size # type: ignore[assignment] + weights["_mlp_size"] = mlp_size # type: ignore[assignment] + + return weights + + +# --------------------------------------------------------------------------- +# 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/python/tensorrt_model_connect/families/smollm3/config.py b/python/tensorrt_model_connect/families/smollm3/config.py new file mode 100644 index 0000000000..164e124ebc --- /dev/null +++ b/python/tensorrt_model_connect/families/smollm3/config.py @@ -0,0 +1,274 @@ +# 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 + + +def _raw_config(config: object) -> dict: + """Return a config object's raw HF mapping. + + Accepts any config the build path may carry. ``engine_builder`` constructs + the shared ``tensorrt_model_connect.config.ModelConfig``, not this module's + dataclass, so anything the builders need from the checkpoint has to be + resolved from ``raw`` rather than from a family-local method. + """ + raw = getattr(config, "raw", None) + if not isinstance(raw, dict): + raise ValueError("SmolLM3 config.raw must be a JSON object") + return raw + + +def resolve_rope_layer_schedule(config: object) -> tuple[bool, ...]: + """Per-layer RoPE flags: ``True`` where the layer applies RoPE. + + SmolLM3 interleaves NoPE layers (no positional encoding) among regular RoPE + layers. The checkpoint publishes this as ``no_rope_layers``, whose entries + are ``1`` where the layer *uses* RoPE and ``0`` where it is a NoPE layer. + When the list is absent it is derived from ``no_rope_layer_interval``: layer + ``i`` is NoPE when ``(i + 1) % interval == 0``, matching the upstream + default of one NoPE layer every four layers. + """ + raw = _raw_config(config) + num_layers = int(getattr(config, "num_hidden_layers", 0) or 0) + if num_layers <= 0: + num_layers = int(raw.get("num_hidden_layers", 0) or 0) + published = raw.get("no_rope_layers") + if published is not None: + if (not isinstance(published, (list, tuple)) + or len(published) < num_layers): + raise ValueError( + "no_rope_layers must be a sequence with at least " + f"num_hidden_layers ({num_layers}) entries, got " + f"{published!r}") + return tuple(bool(int(flag)) for flag in published[:num_layers]) + + interval = raw.get("no_rope_layer_interval") + if interval is None: + return (True,) * num_layers + interval = int(interval) + if interval <= 0: + raise ValueError( + f"no_rope_layer_interval must be positive, got {interval}") + return tuple((idx + 1) % interval != 0 for idx in range(num_layers)) + + +@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 + + def rope_layer_schedule(self) -> tuple[bool, ...]: + """Per-layer RoPE flags for this config. + + Delegates to :func:`resolve_rope_layer_schedule`, which is what the + builders call: they receive the shared ``ModelConfig``, which does not + carry this method. + """ + return resolve_rope_layer_schedule(self) + + @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/smollm3/debug_runner.py b/python/tensorrt_model_connect/families/smollm3/debug_runner.py new file mode 100644 index 0000000000..ebb9d8b73f --- /dev/null +++ b/python/tensorrt_model_connect/families/smollm3/debug_runner.py @@ -0,0 +1,519 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""SmolLM3-owned debug runner adapter.""" + +from __future__ import annotations + +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 - exercised in TRT-free test envs + cudart = None # type: ignore[assignment] + +def _check_cuda(status): + """Raise on CUDA error.""" + if cudart is None: + raise RuntimeError("cuda-python is required for family 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): + """Resolve TRT dtype to a NumPy dtype, including BF16 fallback.""" + 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 family debug_runner execution") + if cudart is None: + raise ImportError("cuda-python is required for family 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 C++ DeviceKvCache behavior exactly. + """ + + 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 + + +def runner_from_bundle( + *, + runtime_strategy: str, + config: dict, + header: dict, + engine_plan: bytes, + bundle_path: str, + distributed_communicator: object | None = None, +) -> object | None: + del bundle_path + if runtime_strategy != "smollm3_decoder_kv_cache": + return None + num_layers = header.get("num_layers", config.get("num_hidden_layers", 1)) + return TrtRunner( + engine_plan=engine_plan, + max_cache_length=header["max_cache_length"], + num_layers=num_layers, + distributed_communicator=distributed_communicator, + ) diff --git a/python/tensorrt_model_connect/families/smollm3/dual_profile_decoder_builder.py b/python/tensorrt_model_connect/families/smollm3/dual_profile_decoder_builder.py new file mode 100644 index 0000000000..d8c76f9cf6 --- /dev/null +++ b/python/tensorrt_model_connect/families/smollm3/dual_profile_decoder_builder.py @@ -0,0 +1,841 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Dual-profile decoder engine builder — single engine, two optimization profiles. + +Produces one TensorRT engine that handles both prefill (multi-token) and +decode (single-token) phases by switching between two optimization profiles +at runtime: + * Profile 0 (prefill): Sq ranges over [1, opt=opt_prefill_length, max=max_prefill_length]. + TensorRT picks batched MHA kernels (e.g. ``_gemm_mha_v2``) at opt Sq. + * Profile 1 (decode): Sq fixed to 1. TensorRT picks the GEMV fast-path + (``_gemv_mha_v1``). + +Both profiles use the same graph and weights — only the optimization +profile differs, so the engine's weights live once in GPU memory and the +C++ runtime creates two ``IExecutionContext``s (one per profile) that +share the engine. + +Scope: covers the same architectural variants the legacy +``standard_decoder_builder`` supports — RMSNorm or LayerNorm; SwiGLU or +GeluFC MLP; RoPE (full / partial / interleaved), learned absolute, or +ALiBi position; sequential or parallel residual; optional q/k_norm, +QKV/output/MLP biases, and a Bloom-style embedding LayerNorm. Quantized +builds (fp8 / int8 ``quant_ctx``) thread Q/DQ insertion through every +projection matmul via ``QuantContext.maybe_quantized_matmul``. Per-layer +debug outputs, hidden-state outputs, and the VL ``embed_input`` path stay +on ``standard_decoder_builder`` for now and are dispatched there from +inside ``build_standard_decoder_engine``. + +Tensor contract for the TensorRT native KV-cache path: + Inputs (Sq varies by profile; cache capacity is static) + token_id int32 (-1,) + position_id int32 (-1,) + cache_write_indices int32 (1,) # update start offset + key_value_lengths int32 (1,) # active length after update + cache_k_i bf16 (1, Hkv, capacity, D) # user-owned static buffer + cache_v_i bf16 (1, Hkv, capacity, D) # user-owned static buffer + Outputs + logits float32 (1, vocab) # last-row sliced inside the engine + present_k_i bf16 (1, Hkv, capacity, D) # aliases cache_k_i + present_v_i bf16 (1, Hkv, capacity, D) # aliases cache_v_i + +The legacy generic path remains available for non-SmolLM3 families that import +this builder and still use dense masks plus incremental present K/V outputs. +""" + +from __future__ import annotations + +import sys +from typing import TYPE_CHECKING + +import numpy as np +from tensorrt_model_connect import trt_compat +from ...native_kv_attention_builder import ( + EXPLICIT_ATTENTION_PREFILL_CHUNK_TOKENS, + add_active_prefix_causal_masks, +) + +from . import graph_ops +from . import graph_blocks +# Runtime import: the schedule is resolved while the graph is built, so this +# cannot live under TYPE_CHECKING the way the annotation-only names do. +from .config import resolve_rope_layer_schedule + +trt = trt_compat.get_trt() + +if TYPE_CHECKING: + from .config import ModelConfig + from .checkpoint_mapper import WeightDict + from ...quantization.context import QuantContext + + +def _const_in_work_dtype( + network: trt.INetworkDefinition, + shape: tuple, + values: np.ndarray, + work_np_dtype: np.dtype, + work_trt_dtype: trt.DataType, +) -> trt.ITensor: + """Create a constant in work_np_dtype storage and cast it to work_trt_dtype. + + Needed for bf16 builds: the dual-profile builder stores bf16 weights + on disk as fp16 (work_np_dtype = np.float16), but the runtime tensor + must be bfloat16 to match the rest of the graph. ``add_constant`` + alone produces an fp16 constant — we need an explicit cast to + bfloat16 so layers like IRotaryEmbeddingLayer (which require all + inputs to share a dtype) accept it. fp16 / fp32 builds are no-ops + because work_np_dtype maps directly to work_trt_dtype. + """ + const = graph_ops.add_constant(network, shape, values, dtype=work_np_dtype) + if const.dtype != work_trt_dtype: + const = network.add_cast(const, work_trt_dtype).get_output(0) + return const + + +def _make_matmul_fn( + network: trt.INetworkDefinition, + dtype: np.dtype, + quant_ctx: "QuantContext | None", +): + """Mirror of ``graph_blocks._make_matmul_fn`` for the dual-profile path. + + Returns a callable ``(lhs, lhs_w, rhs_w, rhs_weights, weight_name) -> ITensor`` + that routes through ``QuantContext.maybe_quantized_matmul`` when present + and falls back to a plain ``add_matmul_rhs_constant`` otherwise. The + ``weight_name`` is the dotted weight key (e.g. ``layer.0.w_q``) used by + the quantization profile to look up scales and the per-layer exclude + pattern. + """ + if quant_ctx is None: + def matmul(lhs, lhs_w, rhs_w, rhs_weights, weight_name): + return graph_ops.add_matmul_rhs_constant( + network, lhs, lhs_w, rhs_w, rhs_weights, dtype=dtype) + return matmul + + def matmul(lhs, lhs_w, rhs_w, rhs_weights, weight_name): + return quant_ctx.maybe_quantized_matmul( + network, lhs, lhs_w, rhs_w, rhs_weights, weight_name, + dtype=dtype) + return matmul + + +def _norm_multi( + network: trt.INetworkDefinition, + inp: trt.ITensor, + hidden: int, + gamma: np.ndarray, + beta: np.ndarray | None, + eps_tensor: trt.ITensor, + norm_type: str, + dtype: np.dtype, +) -> trt.ITensor: + if norm_type == "layernorm": + if beta is None: + beta = np.zeros(hidden, dtype=np.float32) + return graph_ops.add_layer_norm( + network, inp, hidden, gamma, beta, eps_tensor, dtype=dtype) + return graph_ops.add_rms_norm( + network, inp, hidden, gamma, eps_tensor, dtype=dtype) + + +# --------------------------------------------------------------------------- +# MLP helpers. +# --------------------------------------------------------------------------- + + +def _swiglu_mlp( + network: trt.INetworkDefinition, + inp: trt.ITensor, + *, + matmul, + weights: "WeightDict", + prefix: str, + hidden: int, + mlp_size: int, +) -> trt.ITensor: + gate = matmul(inp, hidden, mlp_size, + weights[f"{prefix}.w_gate"], f"{prefix}.w_gate") + up = matmul(inp, hidden, mlp_size, + weights[f"{prefix}.w_up"], f"{prefix}.w_up") + sigmoid = network.add_activation(gate, trt.ActivationType.SIGMOID) + swish = network.add_elementwise( + gate, sigmoid.get_output(0), trt.ElementWiseOperation.PROD) + gated = network.add_elementwise( + swish.get_output(0), up, trt.ElementWiseOperation.PROD) + mlp_out = matmul(gated.get_output(0), mlp_size, hidden, + weights[f"{prefix}.w_down"], f"{prefix}.w_down") + return mlp_out + + +def _gelu_fc_mlp( + network: trt.INetworkDefinition, + inp: trt.ITensor, + *, + matmul, + weights: "WeightDict", + prefix: str, + hidden: int, + mlp_size: int, + activation: str, + work_np_dtype: np.dtype, +) -> trt.ITensor: + fc1 = matmul(inp, hidden, mlp_size, + weights[f"{prefix}.w_fc1"], f"{prefix}.w_fc1") + fc1_bias = weights.get(f"{prefix}.fc1_bias") + if fc1_bias is not None: + fc1 = graph_ops.add_bias_sum(network, fc1, mlp_size, fc1_bias, dtype=work_np_dtype) + activated = graph_ops.add_activation(network, fc1, activation, dtype=work_np_dtype) + fc2 = matmul(activated, mlp_size, hidden, + weights[f"{prefix}.w_fc2"], f"{prefix}.w_fc2") + fc2_bias = weights.get(f"{prefix}.fc2_bias") + if fc2_bias is not None: + fc2 = graph_ops.add_bias_sum(network, fc2, hidden, fc2_bias, dtype=work_np_dtype) + return fc2 + + +# --------------------------------------------------------------------------- +# Config guard. +# --------------------------------------------------------------------------- + + +def _supports_config(config: "ModelConfig", weights: "WeightDict") -> None: + """Reject configs the dual-profile builder cannot handle.""" + model_type = getattr(config, "model_type", "").lower() + if "moe" in model_type or "mamba" in model_type or "rwkv" in model_type: + raise NotImplementedError( + f"dual_profile_decoder_builder does not support model_type={model_type!r}") + if "embedding" not in weights: + raise NotImplementedError("missing embedding weight") + if "final_norm" not in weights: + raise NotImplementedError("missing final_norm weight") + + +# --------------------------------------------------------------------------- +# Main builder. +# --------------------------------------------------------------------------- + + +def build_dual_profile_decoder_engine( + config: "ModelConfig", + weights: "WeightDict", + max_cache_length: int, + *, + precision: str = "fp16", + opt_prefill_length: int = 64, + max_prefill_length: int | None = None, + quant_ctx: "QuantContext | None" = None, + norm_type: str = "rmsnorm", + mlp_type: str = "swiglu", + position_type: str = "rope", + activation: str = "silu", + partial_rotary_factor: float = 1.0, + interleaved_rope: bool = False, + parallel_residual: bool = False, + scale_attn_weights: bool = True, + alibi_bias_scale: float = 1.0, + verbose: bool = False, + dynamic_kv_profile_rows: list[int] | None = None, + profile_mode: str = "dual_profile", + native_kv_cache: bool = False, +) -> bytes: + """Build a prefill/decode-capable dynamic-Sq decoder engine. + + ``norm_type`` / ``mlp_type`` / ``position_type`` / ``activation`` / + ``partial_rotary_factor`` / ``interleaved_rope`` / ``parallel_residual`` / + ``scale_attn_weights`` mirror the same parameters on + ``build_standard_decoder_engine``. + ``alibi_bias_scale`` is multiplied into ALiBi slopes before they are added + through the native attention mask. + + ``quant_ctx`` (optional) routes every projection matmul through + ``QuantContext.maybe_quantized_matmul`` for fp8 / int8 Q/DQ insertion; + when ``None`` the matmuls are plain fp16 / bf16 / fp32. + + ``profile_mode`` controls which optimization profiles are emitted: + + * ``"dual_profile"``: one prefill profile followed by one or more decode + profiles. When ``dynamic_kv_profile_rows`` is provided, the decode side + gets one profile per bucket — letting TriAttention pick the smallest + active KV cache bucket at runtime while still benefitting from batched + prefill on the prompt. + * ``"prefill"``: one prefill profile only. This is used by split-engine + bundles, where decode is served by a separate fixed-Sq=1 engine. + * ``"decode"``: one fixed-Sq=1 profile only. This is the decode half of a + split-engine bundle. + + Dynamic-KV cache bucket profiles are only meaningful in ``dual_profile`` + mode. In either mode, cache_k/cache_v inputs are declared dynamic when + bucket profiles are requested so each profile can constrain their row count. + + ``native_kv_cache`` selects TensorRT's ``IKVCacheUpdateLayer`` and a + primitive attention graph with an explicit active-prefix causal mask. + SmolLM3 enables it internally by default; it is not exposed as a user build + flag. + """ + _supports_config(config, weights) + if profile_mode not in ("dual_profile", "prefill", "decode"): + raise ValueError( + "profile_mode must be 'dual_profile', 'prefill', or 'decode', " + f"got {profile_mode!r}") + + if max_prefill_length is None: + max_prefill_length = max_cache_length + if native_kv_cache: + # Physical KV capacity and one TensorRT enqueue's query length are + # separate limits. Keep the complete model context in the cache while + # bounding the explicit score matrix; the runtime transparently + # advances through multiple chunks. Clamp explicit caller overrides as + # well as the default so they cannot bypass this safety bound. + max_prefill_length = min( + max_prefill_length, + EXPLICIT_ATTENTION_PREFILL_CHUNK_TOKENS, + ) + max_prefill_length = max(1, min(max_prefill_length, max_cache_length)) + opt_prefill_length = max(1, min(opt_prefill_length, max_prefill_length)) + + multi_bucket_decode = bool(dynamic_kv_profile_rows) + if native_kv_cache and multi_bucket_decode: + raise ValueError( + "TensorRT native KV cache requires one fixed physical capacity; " + "dynamic KV-cache bucket profiles are not supported") + if native_kv_cache and position_type == "alibi": + raise NotImplementedError( + "TensorRT native KV cache prototype does not support ALiBi") + if multi_bucket_decode: + decode_buckets: list[int] = [] + seen = set() + for raw in dynamic_kv_profile_rows or []: + clamped = max(1, min(int(raw), max_cache_length)) + if clamped not in seen: + seen.add(clamped) + decode_buckets.append(clamped) + decode_buckets.sort() + if not decode_buckets: + decode_buckets = [max_cache_length] + multi_bucket_decode = False + + attention_size = weights.get("_attention_size", config.attention_size) + mlp_size = weights.get("_mlp_size", config.intermediate_size) + hidden = config.hidden_size + vocab = config.vocab_size + num_layers = config.num_hidden_layers + num_heads = config.num_attention_heads + num_kv_heads = config.num_key_value_heads + head_dim = attention_size // num_heads + kv_attention_size = graph_blocks.infer_kv_attention_size( + weights, num_kv_heads=num_kv_heads, head_dim=head_dim) + rotary_embedding_dim = int(head_dim * partial_rotary_factor) + native_rope_inv_freq: np.ndarray | None = None + if native_kv_cache and position_type == "rope": + rope_scaling = ( + config.raw.get("rope_parameters") + or config.raw.get("rope_scaling") + ) + native_rope_inv_freq = graph_ops.make_native_active_rope_inv_freq( + head_dim, + config.rope_theta, + partial_rotary_factor, + rope_scaling=rope_scaling, + ) + + logger = trt.Logger(trt.Logger.VERBOSE if verbose else trt.Logger.WARNING) + builder = trt.Builder(logger) + network = builder.create_network( + 1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) + trt_config = builder.create_builder_config() + # Native full-context SmolLM3 builds can require substantial tactic workspace + # while TensorRT compiles the primitive attention graph. This is build-time + # scratch only (it is not serialized as runtime KV memory), and the limit + # does not allocate the bytes eagerly. Other paths keep TensorRT's device + # default instead of imposing the former 1 GiB cap. + if native_kv_cache: + trt_config.set_memory_pool_limit( + trt.MemoryPoolType.WORKSPACE, 16 << 30) + + if precision == "fp16": + work_np_dtype, work_trt_dtype = np.float16, trt.float16 + elif precision == "bf16": + work_np_dtype, work_trt_dtype = np.float16, trt.bfloat16 + else: + work_np_dtype, work_trt_dtype = np.float32, trt.float32 + + # ---- Inputs (dynamic Sq) --------------------------------------------- + token_id = network.add_input("token_id", trt.int32, (-1,)) + position_id = network.add_input("position_id", trt.int32, (-1,)) + attention_mask: trt.ITensor | None = None + cache_write_indices: trt.ITensor | None = None + key_value_lengths: trt.ITensor | None = None + if native_kv_cache: + cache_write_indices = network.add_input( + "cache_write_indices", trt.int32, (1,)) + key_value_lengths = network.add_input( + "key_value_lengths", trt.int32, (1,)) + else: + attention_mask = network.add_input( + "attention_mask", trt.float32, (-1, -1)) + + cache_shape: tuple[int, ...] + if native_kv_cache: + cache_shape = (1, num_kv_heads, max_cache_length, head_dim) + elif multi_bucket_decode: + cache_shape = (-1, kv_attention_size) + else: + cache_shape = (max_cache_length, kv_attention_size) + cache_k_inputs: list[trt.ITensor] = [] + cache_v_inputs: list[trt.ITensor] = [] + for i in range(num_layers): + ck = network.add_input( + graph_ops.layer_tensor_name("cache_k", i), + work_trt_dtype, cache_shape) + cv = network.add_input( + graph_ops.layer_tensor_name("cache_v", i), + work_trt_dtype, cache_shape) + cache_k_inputs.append(ck) + cache_v_inputs.append(cv) + + # Cast the legacy dense mask to compute dtype for elementwise broadcast. + attention_mask_work: trt.ITensor | None = attention_mask + if ( + attention_mask is not None + and work_trt_dtype != trt.float32 + ): + attention_mask_work = network.add_cast( + attention_mask, work_trt_dtype).get_output(0) + + # Two (or 1+N) optimization profiles — same graph, different Sq / cache. + def _add_profile(opt_sq: int, max_sq: int, *, fixed: bool = False, + cache_rows_min: int | None = None, + cache_rows_opt: int | None = None, + cache_rows_max: int | None = None): + prof = builder.create_optimization_profile() + min_sq = opt_sq if fixed else 1 + if multi_bucket_decode: + cmn = cache_rows_min if cache_rows_min is not None else max_cache_length + cop = cache_rows_opt if cache_rows_opt is not None else max_cache_length + cmx = cache_rows_max if cache_rows_max is not None else max_cache_length + else: + cmn = cop = cmx = max_cache_length + prof.set_shape("token_id", (min_sq,), (opt_sq,), (max_sq,)) + prof.set_shape("position_id", (min_sq,), (opt_sq,), (max_sq,)) + if not native_kv_cache: + prof.set_shape( + "attention_mask", + (min_sq, cmn + min_sq), + (opt_sq, cop + opt_sq), + (max_sq, cmx + max_sq)) + if multi_bucket_decode: + for i in range(num_layers): + for name in (graph_ops.layer_tensor_name("cache_k", i), + graph_ops.layer_tensor_name("cache_v", i)): + prof.set_shape( + name, + (cmn, kv_attention_size), + (cop, kv_attention_size), + (cmx, kv_attention_size)) + trt_config.add_optimization_profile(prof) + + import os as _os_dbg + if profile_mode == "prefill": + _add_profile(opt_prefill_length, max_prefill_length, fixed=False, + cache_rows_min=1, cache_rows_opt=max_cache_length, + cache_rows_max=max_cache_length) + elif profile_mode == "decode": + _add_profile(1, 1, fixed=True) + elif _os_dbg.environ.get("TRTMC_DECODE_ONLY_DEBUG") == "1": + # Diagnostic: build a one-profile engine with dynamic-shape inputs + # but Sq pinned to 1. Lets us isolate dynamic-shape enqueueV3 + # overhead from per-profile kernel specialisation. + _add_profile(1, 1, fixed=True) + else: + _reverse = _os_dbg.environ.get("TRTMC_REVERSE_PROFILE_ORDER", "0") == "1" + if _reverse: + # Decode profile registered first so it commits its preferred + # weight layout before the prefill profile compiles. + if multi_bucket_decode: + for bucket in decode_buckets: + _add_profile(1, 1, fixed=True, + cache_rows_min=1, cache_rows_opt=bucket, + cache_rows_max=bucket) + else: + _add_profile(1, 1, fixed=True) + _add_profile(opt_prefill_length, max_prefill_length, fixed=False, + cache_rows_min=1, cache_rows_opt=max_cache_length, + cache_rows_max=max_cache_length) + else: + _add_profile(opt_prefill_length, max_prefill_length, fixed=False, + cache_rows_min=1, cache_rows_opt=max_cache_length, + cache_rows_max=max_cache_length) + if multi_bucket_decode: + for bucket in decode_buckets: + _add_profile(1, 1, fixed=True, + cache_rows_min=1, cache_rows_opt=bucket, + cache_rows_max=bucket) + else: + _add_profile(1, 1, fixed=True) + + # ---- Shared constants ------------------------------------------------ + embedding_table = _const_in_work_dtype( + network, (vocab, hidden), weights["embedding"], + work_np_dtype, work_trt_dtype) + + # Native KV derives RoPE only for runtime-active positions, so the engine + # stores a small [D/2] inverse-frequency constant instead of serializing an + # O(context_capacity) table. The legacy graph retains its indexed table. + cos_half_table: trt.ITensor | None = None + sin_half_table: trt.ITensor | None = None + rope_position_id: trt.ITensor | None = position_id + if position_type == "rope": + graph_ops.validate_native_rope_dim(rotary_embedding_dim) + if native_kv_cache: + assert native_rope_inv_freq is not None + cos_half_table, sin_half_table = graph_ops.add_active_rope_cache( + network, + position_id, + native_rope_inv_freq, + work_trt_dtype, + ) + rope_position_id = None + else: + kmax = max_cache_length + max_prefill_length + cos_half_np = graph_ops.make_rope_table_half_dim( + kmax, head_dim, config.rope_theta, True, + partial_rotary_factor, interleaved=interleaved_rope, + rope_scaling=config.raw.get("rope_scaling")) + sin_half_np = graph_ops.make_rope_table_half_dim( + kmax, head_dim, config.rope_theta, False, + partial_rotary_factor, interleaved=interleaved_rope, + rope_scaling=config.raw.get("rope_scaling")) + # BF16 must round directly from the FP32 indexed table. Routing + # through FP16 storage would introduce FP16 -> BF16 double rounding. + rope_np_dtype = ( + np.float32 if work_trt_dtype == trt.bfloat16 else work_np_dtype + ) + cos_half_table = _const_in_work_dtype( + network, cos_half_np.shape, cos_half_np, + rope_np_dtype, work_trt_dtype) + sin_half_table = _const_in_work_dtype( + network, sin_half_np.shape, sin_half_np, + rope_np_dtype, work_trt_dtype) + + # Learned position embedding (GPT-2 / OPT / GPT-Neo / XGLM). + position_embed_table: trt.ITensor | None = None + if position_type == "learned": + pos_embed_np = weights["position_embedding"] + position_embed_table = _const_in_work_dtype( + network, pos_embed_np.shape, pos_embed_np, + work_np_dtype, work_trt_dtype) + + # ALiBi slopes + cache-slot positions for multi-row mask augmentation. + alibi_slopes_tensor: trt.ITensor | None = None + alibi_cache_positions_fp32: trt.ITensor | None = None + if position_type == "alibi": + alibi_slopes_np = graph_ops.compute_alibi_slopes(num_heads) * float(alibi_bias_scale) + # Slopes live as fp32 so the (key_pos - q_pos) math stays in fp32; + # add_alibi_mask_4d casts the final bias to work_trt_dtype before adding + # to the additive mask. + alibi_slopes_tensor = graph_ops.add_constant( + network, (num_heads, 1, 1), + alibi_slopes_np.reshape(num_heads, 1, 1), dtype=np.float32) + # Cache slot k (for k in [0, max_cache_length)) holds the K/V at + # position k. The current step's K/V live in slots + # [max_cache_length, max_cache_length + Sq) and their positions come + # from position_id at runtime, so we only pre-build the cache half. + alibi_cache_positions_fp32 = graph_ops.add_constant( + network, (max_cache_length,), + np.arange(max_cache_length, dtype=np.float32), dtype=np.float32) + + eps_tensor = graph_ops.add_constant( + network, (1, 1), + np.array([[config.rms_norm_eps]], dtype=np.float32), + dtype=np.float32) + eps_tensor_per_head = graph_ops.add_constant( + network, (1, 1, 1), + np.array([[[config.rms_norm_eps]]], dtype=np.float32), + dtype=np.float32) + + # Attention scale. + attn_scale = (1.0 / np.sqrt(max(head_dim, 1))) if scale_attn_weights else 1.0 + + # Quantization-aware matmul (passes weight_name through to QuantContext). + matmul = _make_matmul_fn(network, work_np_dtype, quant_ctx) + + # ---- Embedding ------------------------------------------------------- + emb = network.add_gather(embedding_table, token_id, 0) + hidden_state = emb.get_output(0) # (Sq, hidden) + + if position_type == "learned" and position_embed_table is not None: + pos_gather = network.add_gather(position_embed_table, position_id, 0) + pos_add = network.add_elementwise( + hidden_state, pos_gather.get_output(0), + trt.ElementWiseOperation.SUM) + hidden_state = pos_add.get_output(0) + + # Make sure the main hidden stream is in the requested runtime dtype + # before entering the layer stack (BF16 mode stores fp16 constants). + if hidden_state.dtype != work_trt_dtype: + hidden_state = network.add_cast(hidden_state, work_trt_dtype).get_output(0) + + # Optional embedding LayerNorm (Bloom). + embed_norm = weights.get("embedding_norm") + if embed_norm is not None: + embed_norm_beta = weights.get( + "embedding_norm_beta", np.zeros(hidden, dtype=np.float32)) + hidden_state = _norm_multi( + network, hidden_state, hidden, embed_norm, embed_norm_beta, + eps_tensor, "layernorm", work_np_dtype) + + # The native cache path shares one explicit BOOL mask across every layer. + # The legacy path retains its existing additive-mask graph. + mask_4d: trt.ITensor | None + if native_kv_cache: + mask_4d = None + elif position_type == "alibi": + assert attention_mask_work is not None + mask_4d = graph_ops.add_alibi_mask_4d( + network, attention_mask_work, position_id, + alibi_slopes_tensor, alibi_cache_positions_fp32, + num_heads, target_dtype=work_trt_dtype) + else: + assert attention_mask_work is not None + mask_4d = graph_ops.add_2d_mask_to_4d(network, attention_mask_work) + + native_attention_masks = None + if native_kv_cache: + assert cache_write_indices is not None + assert key_value_lengths is not None + native_attention_masks = add_active_prefix_causal_masks( + network, + token_id, + cache_write_indices, + key_value_lengths, + max_cache_length, + ) + + present_k_outs: list[trt.ITensor] = [] + present_v_outs: list[trt.ITensor] = [] + + # SmolLM3 interleaves NoPE layers; layers flagged False here skip RoPE. + rope_schedule = resolve_rope_layer_schedule(config) + + for layer_idx in range(num_layers): + prefix = f"layer.{layer_idx}" + + # Pre-attention norm. + normed = _norm_multi( + network, hidden_state, hidden, + weights[f"{prefix}.input_norm"], + weights.get(f"{prefix}.input_norm_beta"), + eps_tensor, norm_type, work_np_dtype) + + # Q / K / V projections. + q = matmul(normed, hidden, attention_size, + weights[f"{prefix}.w_q"], f"{prefix}.w_q") + k = matmul(normed, hidden, kv_attention_size, + weights[f"{prefix}.w_k"], f"{prefix}.w_k") + v = matmul(normed, hidden, kv_attention_size, + weights[f"{prefix}.w_v"], f"{prefix}.w_v") + + # Optional QKV biases (Qwen2 / GPT-2 / OPT / Bloom / Falcon / etc.). + q_bias = weights.get(f"{prefix}.q_bias") + if q_bias is not None: + q = graph_ops.add_bias_sum( + network, q, attention_size, q_bias, dtype=work_np_dtype) + k_bias = weights.get(f"{prefix}.k_bias") + if k_bias is not None: + k = graph_ops.add_bias_sum( + network, k, kv_attention_size, k_bias, dtype=work_np_dtype) + v_bias = weights.get(f"{prefix}.v_bias") + if v_bias is not None: + v = graph_ops.add_bias_sum( + network, v, kv_attention_size, v_bias, dtype=work_np_dtype) + + # Optional per-head q/k norm (Qwen3). + q_norm = weights.get(f"{prefix}.q_norm") + if q_norm is not None: + q = graph_ops.add_rms_norm_per_head( + network, q, num_heads, head_dim, q_norm, + eps_tensor_per_head, dtype=work_np_dtype, + sequence_length=None) + k_norm = weights.get(f"{prefix}.k_norm") + if k_norm is not None: + k = graph_ops.add_rms_norm_per_head( + network, k, num_kv_heads, head_dim, k_norm, + eps_tensor_per_head, dtype=work_np_dtype, + sequence_length=None) + + # Position embedding (RoPE only; learned was applied above and ALiBi + # is added into the attention mask). + # NoPE layers carry no positional encoding: q/k stay unrotated. + if position_type == "rope" and rope_schedule[layer_idx]: + q = graph_ops.add_apply_rope_native( + network, q, num_heads, head_dim, + cos_half_table, sin_half_table, rope_position_id, + rotary_embedding_dim, interleaved_rope, + sequence_length=None) + k = graph_ops.add_apply_rope_native( + network, k, num_kv_heads, head_dim, + cos_half_table, sin_half_table, rope_position_id, + rotary_embedding_dim, interleaved_rope, + sequence_length=None) + + if native_kv_cache: + assert cache_write_indices is not None + assert native_attention_masks is not None + native_attention = graph_ops.add_native_kv_cache_attention_from_rows( + network, + q, + k, + v, + cache_k_inputs[layer_idx], + cache_v_inputs[layer_idx], + cache_write_indices, + native_attention_masks, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + q_seq=None, + scale=attn_scale, + tag=f"{prefix}.attn", + ) + context = native_attention["context"] + present_k_outs.append(native_attention["present_k"]) + present_v_outs.append(native_attention["present_v"]) + else: + # Legacy contract: export only this step's raw K/V and materialize + # cache + update with concatenation for attention. + present_k_outs.append(k) + present_v_outs.append(v) + all_k_cat = network.add_concatenation( + [cache_k_inputs[layer_idx], k]) + all_k_cat.axis = 0 + all_v_cat = network.add_concatenation( + [cache_v_inputs[layer_idx], v]) + all_v_cat.axis = 0 + context = graph_ops.add_attention_from_rows( + network, q, all_k_cat.get_output(0), all_v_cat.get_output(0), + num_heads=num_heads, head_dim=head_dim, + num_kv_heads=num_kv_heads, + q_seq=None, kv_seq=None, causal=False, mask=mask_4d, + scale=attn_scale, tag=f"{prefix}.attn") + + attn_out = matmul(context, attention_size, hidden, + weights[f"{prefix}.w_o"], f"{prefix}.w_o") + o_bias = weights.get(f"{prefix}.o_bias") + if o_bias is not None: + attn_out = graph_ops.add_bias_sum( + network, attn_out, hidden, o_bias, dtype=work_np_dtype) + + # Residual structure: parallel (GPT-NeoX / CodeGen / Falcon-3) vs + # sequential (everything else). + if parallel_residual: + post_attn_norm_w = weights.get(f"{prefix}.post_attn_norm") + if post_attn_norm_w is not None: + norm2 = _norm_multi( + network, hidden_state, hidden, + post_attn_norm_w, + weights.get(f"{prefix}.post_attn_norm_beta"), + eps_tensor, norm_type, work_np_dtype) + else: + norm2 = normed + else: + residual1 = network.add_elementwise( + hidden_state, attn_out, trt.ElementWiseOperation.SUM) + norm2 = _norm_multi( + network, residual1.get_output(0), hidden, + weights[f"{prefix}.post_attn_norm"], + weights.get(f"{prefix}.post_attn_norm_beta"), + eps_tensor, norm_type, work_np_dtype) + + # MLP — SwiGLU (Llama-style) or GeluFC (GPT-2-style). + if mlp_type == "gelu_fc": + mlp_out = _gelu_fc_mlp( + network, norm2, + matmul=matmul, weights=weights, prefix=prefix, + hidden=hidden, mlp_size=mlp_size, + activation=activation, work_np_dtype=work_np_dtype) + else: + mlp_out = _swiglu_mlp( + network, norm2, + matmul=matmul, weights=weights, prefix=prefix, + hidden=hidden, mlp_size=mlp_size) + + # Final residual. + if parallel_residual: + sum_attn = network.add_elementwise( + hidden_state, attn_out, trt.ElementWiseOperation.SUM) + residual2 = network.add_elementwise( + sum_attn.get_output(0), mlp_out, trt.ElementWiseOperation.SUM) + else: + residual2 = network.add_elementwise( + residual1.get_output(0), mlp_out, trt.ElementWiseOperation.SUM) + hidden_state = residual2.get_output(0) + + # ---- Final norm + LM head ------------------------------------------- + final_norm = weights.get("final_norm") + if final_norm is not None and len(final_norm) > 0: + hidden_state = _norm_multi( + network, hidden_state, hidden, final_norm, + weights.get("final_norm_beta"), + eps_tensor, norm_type, work_np_dtype) + + # Only the LAST prompt token's logits matter for the next-token sample, + # so slice hidden_state from (Sq, hidden) to (1, hidden) before the LM + # head. This keeps the output contract identical to the single-token + # engine (logits shape = (1, vocab)) under both profiles and avoids + # computing (Sq - 1) redundant vocab-sized matmul rows during prefill. + shape_t = network.add_shape(hidden_state).get_output(0) # [2] int64 + one_hidden = graph_ops.add_constant( + network, (2,), np.array([1, hidden], dtype=np.int64), dtype=np.int64) + start_sub = network.add_elementwise( + shape_t, one_hidden, trt.ElementWiseOperation.SUB) + start_t = start_sub.get_output(0) # [Sq - 1, 0] + size_t = graph_ops.add_constant( + network, (2,), np.array([1, hidden], dtype=np.int64), dtype=np.int64) + slicer = network.add_slice(hidden_state, start=(0, 0), shape=(0, 0), stride=(1, 1)) + slicer.set_input(1, start_t) + slicer.set_input(2, size_t) + last_hidden = slicer.get_output(0) + + out_vocab = (weights["w_out"].shape[1] + if isinstance(weights["w_out"], np.ndarray) else vocab) + logits = graph_ops.add_matmul_rhs_constant( + network, last_hidden, hidden, out_vocab, weights["w_out"], + dtype=work_np_dtype) + lm_bias = weights.get("lm_head_bias") + if lm_bias is not None: + logits = graph_ops.add_bias_sum( + network, logits, out_vocab, lm_bias, dtype=work_np_dtype) + else: + zero_bias = np.zeros(out_vocab, dtype=work_np_dtype) + logits = graph_ops.add_bias_sum( + network, logits, out_vocab, zero_bias, dtype=work_np_dtype) + + if work_trt_dtype != trt.float32: + logits = network.add_cast(logits, trt.float32).get_output(0) + logits.name = "logits" + network.mark_output(logits) + + for i in range(num_layers): + pk = present_k_outs[i] + pv = present_v_outs[i] + pk.name = graph_ops.layer_tensor_name("present_k", i) + pv.name = graph_ops.layer_tensor_name("present_v", i) + network.mark_output(pk) + network.mark_output(pv) + + if verbose: + mode_label = "prefill-profile" if profile_mode == "prefill" else "dual-profile" + print(f"[trtmc build] Building {mode_label} engine " + f"(layers={num_layers}, hidden={hidden}, attn={attention_size}, " + f"kv={kv_attention_size}, " + f"mlp={mlp_size}, cache={max_cache_length}, " + f"opt_prefill={opt_prefill_length}, max_prefill={max_prefill_length}, " + f"norm={norm_type}, mlp_type={mlp_type}, pos={position_type}, " + f"precision={precision}) ...", + file=sys.stderr) + + plan = builder.build_serialized_network(network, trt_config) + if plan is None: + raise RuntimeError("dual-profile decoder engine build failed") + return bytes(plan) diff --git a/python/tensorrt_model_connect/families/smollm3/graph_blocks.py b/python/tensorrt_model_connect/families/smollm3/graph_blocks.py new file mode 100644 index 0000000000..0bbb47fd2f --- /dev/null +++ b/python/tensorrt_model_connect/families/smollm3/graph_blocks.py @@ -0,0 +1,380 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Composable architectural building blocks for TRT engine construction. + +Layer 2 in the three-layer builder stack: + + graph_ops.py Layer 1: Atomic TRT operations (tensor-in/tensor-out) + | + graph_blocks.py Layer 2: Composable blocks (weight-aware) <- THIS FILE + | + builders / plugins Layer 3: Full engine assembly + +Each block composes multiple graph_ops into a reusable sub-structure +(full attention block, SwiGLU MLP, GELU MLP, norm dispatch). Functions +accept a ``weights`` dict + ``prefix`` string to resolve weight names. + +Blocks do NOT apply residual connections. Callers compose the residual +pattern, which is what varies across architectures (sequential vs parallel +residual, DeepStack injection, MoE routing, etc.). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +from tensorrt_model_connect import trt_compat + +from . import graph_ops + +trt = trt_compat.get_trt() + +if TYPE_CHECKING: + from .checkpoint_mapper import WeightDict + from ...quantization.context import QuantContext + + +# --------------------------------------------------------------------------- +# Precision boundary helpers (used by standard_decoder_builder, not inside +# blocks themselves). +# --------------------------------------------------------------------------- + +def make_matmul_fn(network, dtype, quant_ctx): + """Create a matmul callable that routes through quant_ctx if present. + + Returns a function: (lhs, lhs_w, rhs_w, rhs_weights, weight_name) -> ITensor + """ + if quant_ctx is None: + def matmul(lhs, lhs_w, rhs_w, rhs_weights, weight_name): + return graph_ops.add_matmul_rhs_constant( + network, lhs, lhs_w, rhs_w, rhs_weights, dtype=dtype) + return matmul + else: + def matmul(lhs, lhs_w, rhs_w, rhs_weights, weight_name): + return quant_ctx.maybe_quantized_matmul( + network, lhs, lhs_w, rhs_w, rhs_weights, weight_name, + dtype=dtype) + return matmul + + +_make_matmul_fn = make_matmul_fn + + +def infer_kv_attention_size( + weights: dict, + *, + prefix: str = "layer.0", + num_kv_heads: int, + head_dim: int, +) -> int: + """Validate and return the compact K/V row width.""" + expected = int(num_kv_heads * head_dim) + explicit = weights.get("_kv_attention_size") + if explicit is not None and int(explicit) != expected: + raise ValueError( + f"Compact K/V cache width must be num_kv_heads * head_dim " + f"({expected}), got _kv_attention_size={int(explicit)}") + w_k = weights.get(f"{prefix}.w_k") + if isinstance(w_k, np.ndarray) and w_k.ndim == 2: + actual = int(w_k.shape[1]) + if actual != expected: + raise ValueError( + f"{prefix}.w_k must use compact K/V width {expected}, " + f"got {actual}") + return expected + + +def apply_norm( + network: trt.INetworkDefinition, + inp: trt.ITensor, + hidden_size: int, + gamma: np.ndarray, + beta: np.ndarray | None, + eps_tensor: trt.ITensor, + norm_type: str, + dtype: np.dtype = np.float32, + eps: float | None = None, +) -> trt.ITensor: + """Dispatch to RMSNorm or LayerNorm based on norm_type.""" + if norm_type == "layernorm": + if beta is None: + beta = np.zeros(hidden_size, dtype=np.float32) + if eps is not None: + return graph_ops.add_layer_norm_native( + network, inp, hidden_size, gamma, beta, eps, dtype=dtype) + # Native INormalizationLayer requires a build-time scalar epsilon. + # Some callers only pass epsilon as an ITensor, so keep the manual + # shared fallback until those builders thread the scalar too. + return graph_ops.add_layer_norm( + network, inp, hidden_size, gamma, beta, eps_tensor, dtype=dtype) + else: + return graph_ops.add_rms_norm( + network, inp, hidden_size, gamma, eps_tensor, dtype=dtype) + + +def add_attention_block( + network: trt.INetworkDefinition, + hidden: trt.ITensor, + cache_k: trt.ITensor, + cache_v: trt.ITensor, + attention_mask: trt.ITensor, + position_id: trt.ITensor, + *, + weights: WeightDict, + prefix: str, + hidden_size: int, + attention_size: int, + num_heads: int, + head_dim: int, + max_cache_length: int, + eps_tensor: trt.ITensor, + kv_attention_size: int | None = None, + num_kv_heads: int | None = None, + attention_scale: float | None = None, + eps: float | None = None, + norm_type: str = "rmsnorm", + position_type: str = "rope", + apply_rope: bool = True, + alibi_slopes_tensor: trt.ITensor | None = None, + alibi_indices_tensor: trt.ITensor | None = None, + dtype: np.dtype = np.float32, + quant_ctx: QuantContext | None = None, + layer_prefix: str = "", + # TRT 10 native API tensors. + cos_half_tensor: trt.ITensor | None = None, + sin_half_tensor: trt.ITensor | None = None, + rotary_embedding_dim: int = 0, + interleaved_rope: bool = False, + ffi_attention_kernel: str | None = None, + dynamic_kv_cache: bool = False, +) -> dict[str, trt.ITensor]: + """Pre-norm -> QKV -> RoPE -> cache concat -> attention -> output proj. + + Returns {"normed": ..., "attn_out": ..., "present_k": ..., "present_v": ...}. + Does NOT apply residual -- callers compose the residual pattern. + + This function uses TRT 10 native APIs for the basic transformer primitives: + - IRotaryEmbeddingLayer for RoPE + - IAttention for scaled dot-product attention + ALiBi is represented as a per-head additive attention mask and still uses + native IAttention. + """ + matmul = _make_matmul_fn(network, dtype, quant_ctx) + attention_window = max_cache_length + 1 + if num_kv_heads is None: + num_kv_heads = num_heads + if kv_attention_size is None: + kv_attention_size = num_kv_heads * head_dim + + # Weight name for quant scale lookup — use layer_prefix if provided, + # otherwise fall back to the weights-dict prefix. + _lp = layer_prefix or prefix + + # Pre-attention norm + normed = apply_norm( + network, hidden, hidden_size, + weights[f"{prefix}.input_norm"], + weights.get(f"{prefix}.input_norm_beta"), + eps_tensor, norm_type, dtype=dtype, eps=eps) + + # QKV projections + q = matmul(normed, hidden_size, attention_size, + weights[f"{prefix}.w_q"], f"{_lp}.w_q") + k = matmul(normed, hidden_size, kv_attention_size, + weights[f"{prefix}.w_k"], f"{_lp}.w_k") + v = matmul(normed, hidden_size, kv_attention_size, + weights[f"{prefix}.w_v"], f"{_lp}.w_v") + + # Optional QKV biases + q_bias = weights.get(f"{prefix}.q_bias") + if q_bias is not None: + q = graph_ops.add_bias_sum(network, q, attention_size, q_bias, dtype=dtype) + k_bias = weights.get(f"{prefix}.k_bias") + if k_bias is not None: + k = graph_ops.add_bias_sum(network, k, kv_attention_size, k_bias, dtype=dtype) + v_bias = weights.get(f"{prefix}.v_bias") + if v_bias is not None: + v = graph_ops.add_bias_sum(network, v, kv_attention_size, v_bias, dtype=dtype) + + # Optional per-head q/k norm + q_norm = weights.get(f"{prefix}.q_norm") + if q_norm is not None: + q = graph_ops.add_rms_norm_per_head( + network, q, num_heads, head_dim, q_norm, eps_tensor, dtype=dtype) + k_norm = weights.get(f"{prefix}.k_norm") + if k_norm is not None: + k = graph_ops.add_rms_norm_per_head( + network, k, num_kv_heads, head_dim, k_norm, eps_tensor, dtype=dtype) + + # ------------------------------------------------------------------ # + # RoPE via native IRotaryEmbeddingLayer # + # ------------------------------------------------------------------ # + use_native_attention = ffi_attention_kernel is None + + # ``apply_rope`` is False on NoPE layers, which carry no positional + # encoding at all: q/k reach attention unrotated. + if position_type == "rope" and apply_rope: + if cos_half_tensor is None or sin_half_tensor is None: + raise ValueError( + "RoPE attention requires half-dimension cos/sin tensors for " + "TRT native IRotaryEmbeddingLayer") + rope_dim = rotary_embedding_dim or head_dim + rope_dim = graph_ops.validate_native_rope_dim(rope_dim) + q = graph_ops.add_apply_rope_native( + network, q, num_heads, head_dim, + cos_half_tensor, sin_half_tensor, position_id, + rope_dim, interleaved_rope) + k = graph_ops.add_apply_rope_native( + network, k, num_kv_heads, head_dim, + cos_half_tensor, sin_half_tensor, position_id, + rope_dim, interleaved_rope) + + # Save present K/V (before concatenation, this is the raw projection output) + present_k = k + present_v = v + + # Reshape current K, V for concatenation + k_reshape = network.add_shuffle(k) + k_reshape.reshape_dims = (1, kv_attention_size) + v_reshape = network.add_shuffle(v) + v_reshape.reshape_dims = (1, kv_attention_size) + + # Concatenate with cache + all_k = network.add_concatenation( + [cache_k, k_reshape.get_output(0)]) + all_k.axis = 0 + all_v = network.add_concatenation( + [cache_v, v_reshape.get_output(0)]) + all_v.axis = 0 + + # ------------------------------------------------------------------ # + # Attention core — native IAttention or FFI kernel # + # ------------------------------------------------------------------ # + if use_native_attention: + kv_seq = None if dynamic_kv_cache else attention_window + if alibi_slopes_tensor is not None: + if alibi_indices_tensor is None: + raise ValueError("ALiBi attention requires cache position indices") + if dynamic_kv_cache: + raise ValueError("dynamic_kv_cache is not supported for ALiBi attention") + mask_4d = graph_ops.add_alibi_mask_4d( + network, + attention_mask, + position_id, + alibi_slopes_tensor, + alibi_indices_tensor, + num_heads, + ) + else: + mask_4d = graph_ops.add_2d_mask_to_4d(network, attention_mask) + + context = graph_ops.add_attention_from_rows( + network, + q, + all_k.get_output(0), + all_v.get_output(0), + num_heads=num_heads, + head_dim=head_dim, + num_kv_heads=num_kv_heads, + q_seq=1, + kv_seq=kv_seq, + causal=False, + mask=mask_4d, + scale=attention_scale, + ) + elif ffi_attention_kernel is not None: + if num_kv_heads != num_heads: + raise ValueError( + "FFI decoder attention requires num_kv_heads == num_heads; " + "use TRT native attention for compact GQA/MQA KV cache") + # Fused attention kernel via TVM-FFI plugin + context = graph_ops.add_decoder_attention_ffi( + network, q, all_k.get_output(0), all_v.get_output(0), + kernel_name=ffi_attention_kernel, + num_heads=num_heads, head_dim=head_dim, + attention_window=attention_window) + + # Output projection + attn_out = matmul(context, + attention_size, hidden_size, + weights[f"{prefix}.w_o"], f"{_lp}.w_o") + + # Optional output projection bias + o_bias = weights.get(f"{prefix}.o_bias") + if o_bias is not None: + attn_out = graph_ops.add_bias_sum(network, attn_out, hidden_size, o_bias, dtype=dtype) + + return { + "normed": normed, + "attn_out": attn_out, + "present_k": present_k, + "present_v": present_v, + } + + +def add_swiglu_mlp( + network: trt.INetworkDefinition, + inp: trt.ITensor, + *, + weights: WeightDict, + prefix: str, + hidden_size: int, + mlp_size: int, + dtype: np.dtype = np.float32, + quant_ctx: QuantContext | None = None, + layer_prefix: str = "", +) -> trt.ITensor: + """Gate/up/down SwiGLU MLP. Returns output tensor.""" + matmul = _make_matmul_fn(network, dtype, quant_ctx) + _lp = layer_prefix or prefix + + gate = matmul(inp, hidden_size, mlp_size, + weights[f"{prefix}.w_gate"], f"{_lp}.w_gate") + up = matmul(inp, hidden_size, mlp_size, + weights[f"{prefix}.w_up"], f"{_lp}.w_up") + + sigmoid = network.add_activation(gate, trt.ActivationType.SIGMOID) + swish = network.add_elementwise( + gate, sigmoid.get_output(0), trt.ElementWiseOperation.PROD) + gated = network.add_elementwise( + swish.get_output(0), up, trt.ElementWiseOperation.PROD) + + mlp_out = matmul(gated.get_output(0), mlp_size, hidden_size, + weights[f"{prefix}.w_down"], f"{_lp}.w_down") + return mlp_out + + +def add_gelu_fc_mlp( + network: trt.INetworkDefinition, + inp: trt.ITensor, + *, + weights: WeightDict, + prefix: str, + hidden_size: int, + mlp_size: int, + activation: str = "gelu_new", + dtype: np.dtype = np.float32, + quant_ctx: QuantContext | None = None, + layer_prefix: str = "", +) -> trt.ITensor: + """fc1 -> activation -> fc2 MLP. Returns output tensor.""" + matmul = _make_matmul_fn(network, dtype, quant_ctx) + _lp = layer_prefix or prefix + + fc1 = matmul(inp, hidden_size, mlp_size, + weights[f"{prefix}.w_fc1"], f"{_lp}.w_fc1") + fc1_bias = weights.get(f"{prefix}.fc1_bias") + if fc1_bias is not None: + fc1 = graph_ops.add_bias_sum(network, fc1, mlp_size, fc1_bias, dtype=dtype) + + activated = graph_ops.add_activation(network, fc1, activation, dtype=dtype) + + fc2 = matmul(activated, mlp_size, hidden_size, + weights[f"{prefix}.w_fc2"], f"{_lp}.w_fc2") + fc2_bias = weights.get(f"{prefix}.fc2_bias") + if fc2_bias is not None: + fc2 = graph_ops.add_bias_sum(network, fc2, hidden_size, fc2_bias, dtype=dtype) + + return fc2 diff --git a/python/tensorrt_model_connect/families/smollm3/graph_ops.py b/python/tensorrt_model_connect/families/smollm3/graph_ops.py new file mode 100644 index 0000000000..c5ce1853c2 --- /dev/null +++ b/python/tensorrt_model_connect/families/smollm3/graph_ops.py @@ -0,0 +1,1444 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Family-owned TensorRT graph operations for Python engine builds. + +Tensor names and shapes must stay compatible with the C++ bundle runtime. +""" + +from __future__ import annotations + +from collections.abc import Mapping +import math +from typing import Any + +import numpy as np +from tensorrt_model_connect import trt_compat +from ...native_kv_attention_builder import ( + NativeKvMasks, + add_explicit_masked_grouped_query_attention, +) + + +trt = trt_compat.get_trt() + + +def _cast_back_to_trt_dtype( + network: trt.INetworkDefinition, + tensor: trt.ITensor, + target_dtype: trt.DataType, +) -> trt.ITensor: + """Cast a tensor back to the original TRT runtime dtype after FP32 compute.""" + if tensor.dtype == target_dtype: + return tensor + return network.add_cast(tensor, target_dtype).get_output(0) + + +def _add_matrix_multiply_with_fp32_accumulation( + network: trt.INetworkDefinition, + lhs: trt.ITensor, + lhs_op: trt.MatrixOperation, + rhs: trt.ITensor, + rhs_op: trt.MatrixOperation, +) -> trt.ITensor: + """Request TensorRT's fused FP16 GEMM with FP32 accumulation.""" + output_dtype = lhs.dtype + if lhs.dtype == trt.float16 and rhs.dtype == trt.float16: + lhs = network.add_cast(lhs, trt.float32).get_output(0) + rhs = network.add_cast(rhs, trt.float32).get_output(0) + output = network.add_matrix_multiply(lhs, lhs_op, rhs, rhs_op).get_output(0) + return _cast_back_to_trt_dtype(network, output, output_dtype) + +def layer_tensor_name(stem: str, layer: int) -> str: + return f"{stem}_{layer}" + + +def add_constant( + network: trt.INetworkDefinition, + shape: tuple[int, ...], + values: np.ndarray, + dtype: np.dtype = np.float32, +) -> trt.ITensor: + """Add a constant tensor in the given *dtype* (default float32).""" + weights = trt.Weights(np.ascontiguousarray(values, dtype=dtype)) + layer = network.add_constant(shape, weights) + return layer.get_output(0) + + +def add_matmul_rhs_constant( + network: trt.INetworkDefinition, + lhs: trt.ITensor, + lhs_width: int, + rhs_width: int, + rhs_weights: np.ndarray, + dtype: np.dtype = np.float32, + fp32_accumulation: bool = True, +) -> trt.ITensor: + """Matrix multiply: lhs @ rhs_constant. rhs is [lhs_width, rhs_width].""" + rank = len(tuple(lhs.shape)) + rhs_shape = ( + (lhs_width, rhs_width) + if rank <= 2 + else (1,) * (rank - 2) + (lhs_width, rhs_width) + ) + rhs = add_constant( + network, + rhs_shape, + np.asarray(rhs_weights).reshape(rhs_shape), + dtype=dtype, + ) + rhs = _cast_back_to_trt_dtype(network, rhs, lhs.dtype) + if fp32_accumulation: + return _add_matrix_multiply_with_fp32_accumulation( + network, + lhs, trt.MatrixOperation.NONE, + rhs, trt.MatrixOperation.NONE, + ) + mm = network.add_matrix_multiply( + lhs, trt.MatrixOperation.NONE, + rhs, trt.MatrixOperation.NONE, + ) + return _cast_back_to_trt_dtype(network, mm.get_output(0), lhs.dtype) + + +def add_bias_sum( + network: trt.INetworkDefinition, + inp: trt.ITensor, + width: int, + bias: np.ndarray, + dtype: np.dtype = np.float32, +) -> trt.ITensor: + """Element-wise add a bias broadcast over all non-feature axes.""" + rank = len(tuple(inp.shape)) + bias_shape = (width,) if rank <= 1 else (1,) * (rank - 1) + (width,) + bias_t = add_constant( + network, bias_shape, np.asarray(bias).reshape(bias_shape), dtype=dtype) + bias_t = _cast_back_to_trt_dtype(network, bias_t, inp.dtype) + s = network.add_elementwise(inp, bias_t, trt.ElementWiseOperation.SUM) + return _cast_back_to_trt_dtype(network, s.get_output(0), inp.dtype) + + +def add_rms_norm( + network: trt.INetworkDefinition, + inp: trt.ITensor, + hidden_size: int, + gamma: np.ndarray, + eps_tensor: trt.ITensor, + dtype: np.dtype = np.float32, +) -> trt.ITensor: + """RMSNorm: gamma * (x / sqrt(mean(x^2) + eps)). + + FP32 precision boundary: when dtype != float32, casts to FP32 before + norm computation for numerical stability, then casts back. + + TRT's native normalization API implements mean-centered LayerNorm, not + RMSNorm, so this remains a manual shared implementation. + """ + need_cast = (dtype != np.float32) + output_dtype = inp.dtype + if need_cast: + inp = network.add_cast(inp, trt.float32).get_output(0) + eps_tensor = network.add_cast(eps_tensor, trt.float32).get_output(0) + sq = network.add_elementwise(inp, inp, trt.ElementWiseOperation.PROD) + mean = network.add_reduce( + sq.get_output(0), trt.ReduceOperation.AVG, 1 << 1, keep_dims=True) + denom_in = network.add_elementwise( + mean.get_output(0), eps_tensor, trt.ElementWiseOperation.SUM) + sqrt_l = network.add_unary(denom_in.get_output(0), trt.UnaryOperation.SQRT) + recip = network.add_unary(sqrt_l.get_output(0), trt.UnaryOperation.RECIP) + normalized = network.add_elementwise( + inp, recip.get_output(0), trt.ElementWiseOperation.PROD) + gamma_t = add_constant(network, (1, hidden_size), gamma, dtype=np.float32) + scaled = network.add_elementwise( + normalized.get_output(0), gamma_t, trt.ElementWiseOperation.PROD) + result = scaled.get_output(0) + if need_cast: + result = _cast_back_to_trt_dtype(network, result, output_dtype) + return result + + +def add_rms_norm_per_head( + network: trt.INetworkDefinition, + inp: trt.ITensor, + num_heads: int, + head_dim: int, + gamma: np.ndarray, + eps_tensor: trt.ITensor, + dtype: np.dtype = np.float32, + sequence_length: int | None = 1, +) -> trt.ITensor: + """Per-head RMSNorm for [Sq, num_heads * head_dim] tensors. + + FP32 precision boundary: when dtype != float32, casts to FP32 before + norm computation for numerical stability, then casts back. + ``sequence_length=None`` means runtime-dynamic Sq. + ``gamma`` may be [num_heads * head_dim] or [head_dim] broadcast to heads. + """ + need_cast = (dtype != np.float32) + output_dtype = inp.dtype + seq_dim = -1 if sequence_length is None else sequence_length + reshape_in = network.add_shuffle(inp) + reshape_in.reshape_dims = (seq_dim, num_heads, head_dim) + + reshaped = reshape_in.get_output(0) + if need_cast: + reshaped = network.add_cast(reshaped, trt.float32).get_output(0) + eps_tensor = network.add_cast(eps_tensor, trt.float32).get_output(0) + eps_3d = network.add_shuffle(eps_tensor) + eps_3d.reshape_dims = (1, 1, 1) + sq = network.add_elementwise(reshaped, reshaped, trt.ElementWiseOperation.PROD) + mean = network.add_reduce( + sq.get_output(0), trt.ReduceOperation.AVG, 1 << 2, keep_dims=True) + denom_in = network.add_elementwise( + mean.get_output(0), eps_3d.get_output(0), trt.ElementWiseOperation.SUM) + sqrt_l = network.add_unary(denom_in.get_output(0), trt.UnaryOperation.SQRT) + recip = network.add_unary(sqrt_l.get_output(0), trt.UnaryOperation.RECIP) + normalized = network.add_elementwise( + reshaped, recip.get_output(0), trt.ElementWiseOperation.PROD) + gamma_arr = np.asarray(gamma, dtype=np.float32) + if gamma_arr.size == head_dim: + gamma_t = add_constant( + network, (1, 1, head_dim), gamma_arr.reshape(1, 1, head_dim), + dtype=np.float32) + else: + gamma_t = add_constant( + network, (1, num_heads, head_dim), + gamma_arr.reshape(num_heads, head_dim), dtype=np.float32) + scaled = network.add_elementwise( + normalized.get_output(0), gamma_t, trt.ElementWiseOperation.PROD) + + result = scaled.get_output(0) + if need_cast: + result = _cast_back_to_trt_dtype(network, result, output_dtype) + reshape_out = network.add_shuffle(result) + reshape_out.reshape_dims = (seq_dim, num_heads * head_dim) + return reshape_out.get_output(0) + + +def add_layer_norm( + network: trt.INetworkDefinition, + inp: trt.ITensor, + hidden_size: int, + gamma: np.ndarray, + beta: np.ndarray, + eps_tensor: trt.ITensor, + dtype: np.dtype = np.float32, +) -> trt.ITensor: + """LayerNorm: gamma * ((x - mean) / sqrt(var + eps)) + beta. + + FP32 precision boundary: when dtype != float32, casts to FP32 before + norm computation for numerical stability, then casts back. + """ + need_cast = (dtype != np.float32) + output_dtype = inp.dtype + if need_cast: + inp = network.add_cast(inp, trt.float32).get_output(0) + eps_tensor = network.add_cast(eps_tensor, trt.float32).get_output(0) + # mean = reduce_mean(x) + mean = network.add_reduce( + inp, trt.ReduceOperation.AVG, 1 << 1, keep_dims=True) + # x - mean + centered = network.add_elementwise( + inp, mean.get_output(0), trt.ElementWiseOperation.SUB) + # variance = mean((x - mean)^2) + sq = network.add_elementwise( + centered.get_output(0), centered.get_output(0), + trt.ElementWiseOperation.PROD) + var = network.add_reduce( + sq.get_output(0), trt.ReduceOperation.AVG, 1 << 1, keep_dims=True) + # sqrt(var + eps) + denom_in = network.add_elementwise( + var.get_output(0), eps_tensor, trt.ElementWiseOperation.SUM) + sqrt_l = network.add_unary(denom_in.get_output(0), trt.UnaryOperation.SQRT) + recip = network.add_unary(sqrt_l.get_output(0), trt.UnaryOperation.RECIP) + # normalized = (x - mean) / sqrt(var + eps) + normalized = network.add_elementwise( + centered.get_output(0), recip.get_output(0), + trt.ElementWiseOperation.PROD) + # gamma * normalized + beta + gamma_t = add_constant(network, (1, hidden_size), gamma, dtype=np.float32) + scaled = network.add_elementwise( + normalized.get_output(0), gamma_t, trt.ElementWiseOperation.PROD) + beta_t = add_constant(network, (1, hidden_size), beta, dtype=np.float32) + result = network.add_elementwise( + scaled.get_output(0), beta_t, trt.ElementWiseOperation.SUM) + result = result.get_output(0) + if need_cast: + result = _cast_back_to_trt_dtype(network, result, output_dtype) + return result + + +def add_gelu_new( + network: trt.INetworkDefinition, + inp: trt.ITensor, + dtype: np.dtype = np.float32, +) -> trt.ITensor: + """GELU (tanh approximation): 0.5*x*(1+tanh(sqrt(2/pi)*(x+0.044715*x^3))). + + Constants are cast to ``inp.dtype`` so the elementwise ops are valid in + a STRONGLY_TYPED network when ``inp`` is bf16 (storage np_dtype is + fp16, runtime trt_dtype is bfloat16) or any other non-matching combo. + """ + target_dtype = inp.dtype + const_shape = (1,) * max(1, len(tuple(inp.shape))) + + def _const(name, value): + c = add_constant( + network, const_shape, np.array([value], dtype=np.float32), dtype=dtype) + return _cast_back_to_trt_dtype(network, c, target_dtype) + + # x^3 + x_sq = network.add_elementwise(inp, inp, trt.ElementWiseOperation.PROD) + x_cu = network.add_elementwise( + x_sq.get_output(0), inp, trt.ElementWiseOperation.PROD) + # 0.044715 * x^3 + coeff = _const("coeff", 0.044715) + scaled_cube = network.add_elementwise( + x_cu.get_output(0), coeff, trt.ElementWiseOperation.PROD) + # x + 0.044715 * x^3 + inner_sum = network.add_elementwise( + inp, scaled_cube.get_output(0), trt.ElementWiseOperation.SUM) + # sqrt(2/pi) * (x + 0.044715 * x^3) + sqrt_2_over_pi = _const("sqrt_2_over_pi", np.sqrt(2.0 / np.pi)) + tanh_arg = network.add_elementwise( + sqrt_2_over_pi, inner_sum.get_output(0), + trt.ElementWiseOperation.PROD) + # tanh(...) + tanh_l = network.add_activation( + tanh_arg.get_output(0), trt.ActivationType.TANH) + # 1 + tanh(...) + one = _const("one", 1.0) + one_plus_tanh = network.add_elementwise( + one, tanh_l.get_output(0), trt.ElementWiseOperation.SUM) + # 0.5 * x + half = _const("half", 0.5) + half_x = network.add_elementwise( + half, inp, trt.ElementWiseOperation.PROD) + # 0.5 * x * (1 + tanh(...)) + result = network.add_elementwise( + half_x.get_output(0), one_plus_tanh.get_output(0), + trt.ElementWiseOperation.PROD) + return result.get_output(0) + + +def add_activation( + network: trt.INetworkDefinition, + inp: trt.ITensor, + activation_type: str, + dtype: np.dtype = np.float32, +) -> trt.ITensor: + """Dispatch activation by name: 'silu', 'gelu_new', 'gelu', 'relu', 'relu2'/'squared_relu'.""" + if activation_type in ("gelu_new", "gelu"): + return add_gelu_new(network, inp, dtype=dtype) + elif activation_type == "relu": + act = network.add_activation(inp, trt.ActivationType.RELU) + return act.get_output(0) + elif activation_type in ("relu2", "squared_relu"): + relu = network.add_activation(inp, trt.ActivationType.RELU) + sq = network.add_elementwise( + relu.get_output(0), relu.get_output(0), + trt.ElementWiseOperation.PROD) + return sq.get_output(0) + elif activation_type == "silu": + sigmoid = network.add_activation(inp, trt.ActivationType.SIGMOID) + swish = network.add_elementwise( + inp, sigmoid.get_output(0), trt.ElementWiseOperation.PROD) + return swish.get_output(0) + else: + raise ValueError(f"Unsupported activation: {activation_type}") + + +def compute_alibi_slopes(num_heads: int) -> np.ndarray: + """Compute ALiBi slopes for each attention head (from the ALiBi paper). + + For power-of-2 num_heads: geometric sequence 2^(-8/n * i), i in 1..n. + For non-power-of-2: interleave two geometric sequences. + + Returns: [num_heads] float32 array. + """ + def _get_slopes_power_of_2(n: int) -> list[float]: + start = 2 ** (-(2 ** -(np.log2(n) - 3))) + return [start * (start ** i) for i in range(n)] + + if num_heads > 0 and (num_heads & (num_heads - 1)) == 0: + # Power of 2 + return np.array(_get_slopes_power_of_2(num_heads), dtype=np.float32) + else: + closest_power_of_2 = 2 ** int(np.floor(np.log2(num_heads))) + slopes_a = _get_slopes_power_of_2(closest_power_of_2) + slopes_b = _get_slopes_power_of_2(2 * closest_power_of_2) + slopes_b = slopes_b[0::2][: num_heads - closest_power_of_2] + return np.array(slopes_a + slopes_b, dtype=np.float32) + + +# Alias: add_gelu_tanh is the same as add_gelu_new (tanh approximation) +add_gelu_tanh = add_gelu_new + + +# --------------------------------------------------------------------------- +# TRT 10 native attention APIs (TRT 10.x) +# +# Three primitives replace manual primitive chains: +# add_layer_norm_native → INormalizationLayer (replaces add_layer_norm) +# add_apply_rope_native → IRotaryEmbeddingLayer +# add_attention_core → IAttention (replaces score+softmax+V) +# --------------------------------------------------------------------------- + +def add_layer_norm_native( + network: trt.INetworkDefinition, + inp: trt.ITensor, + hidden_size: int, + gamma: np.ndarray, + beta: np.ndarray, + eps: float, + dtype: np.dtype = np.float32, +) -> trt.ITensor: + """LayerNorm via TRT native INormalizationLayer (add_normalization_v2). + + Replaces the manual reduce/elementwise chain in add_layer_norm with a + single fused layer that TRT can optimize end-to-end. In strongly typed + networks, input/scale/bias must have identical tensor types; compute + precision is set to FP32 for numerical stability when the TensorRT Python + layer exposes that control. + + Note: INormalizationLayer computes (x - mean) / sqrt(var + eps) * gamma + beta. + This is LayerNorm, NOT RMSNorm. Use add_rms_norm for RMSNorm models. + + Args: + inp: Input tensor [*, hidden_size]. + hidden_size: Size of the normalized dimension (last axis). + gamma: Scale weights [hidden_size]. + beta: Bias weights [hidden_size]. + eps: Numerical stability epsilon (scalar, not a tensor). + dtype: Storage dtype for gamma/beta constants before TRT cast. + """ + inp_shape = getattr(inp, "shape", None) + rank = len(tuple(inp_shape)) if inp_shape is not None else 2 + param_shape = ( + (hidden_size,) if rank <= 1 else (1,) * (rank - 1) + (hidden_size,) + ) + gamma_t = add_constant( + network, param_shape, np.asarray(gamma).reshape(param_shape), dtype=dtype) + beta_t = add_constant( + network, param_shape, np.asarray(beta).reshape(param_shape), dtype=dtype) + gamma_t = _cast_back_to_trt_dtype(network, gamma_t, inp.dtype) + beta_t = _cast_back_to_trt_dtype(network, beta_t, inp.dtype) + # axesMask bit i selects axis i as a reduction axis. The normalized + # hidden dimension is always the last axis for [*, hidden_size] tensors. + norm = network.add_normalization_v2(inp, gamma_t, beta_t, 1 << (rank - 1)) + norm.epsilon = eps + # TensorRT 11 removed the Python INormalizationLayer.compute_precision + # attribute. Keep the TRT 10 hint, and let TRT 11 infer the precision. + if hasattr(norm, "compute_precision"): + norm.compute_precision = trt.float32 + return norm.get_output(0) + + +def validate_native_rope_dim( + rotary_embedding_dim: int, + *, + field_name: str = "rotary_embedding_dim", +) -> int: + """Validate the dimension contract required by TRT native RoPE.""" + rotary_embedding_dim = int(rotary_embedding_dim) + if rotary_embedding_dim < 2 or rotary_embedding_dim % 2 != 0: + raise ValueError( + f"TRT native RoPE requires {field_name} to be an even value >= 2; " + f"got {rotary_embedding_dim}") + return rotary_embedding_dim + + +def make_native_active_rope_inv_freq( + head_dim: int, + rope_theta: float, + partial_rotary_factor: float = 1.0, + *, + rope_scaling: Mapping[str, Any] | None = None, +) -> np.ndarray: + """Return HF/Torch-exact frequencies for native active-position RoPE. + + Hugging Face initializes Llama inverse frequencies on CPU with Torch before + moving the model to the inference device. NumPy ``power`` can differ from + that operation by one ULP at long positions. The legacy indexed-table + helper below is intentionally unchanged. + """ + rotary_ndims = validate_native_rope_dim( + int(head_dim * partial_rotary_factor)) + rope_theta = float(rope_theta) + if not np.isfinite(rope_theta) or rope_theta <= 0.0: + raise ValueError( + "TRT native RoPE requires rope_theta to be finite and positive; " + f"got {rope_theta}") + try: + import torch + except (ImportError, OSError) as exc: + raise RuntimeError( + "TensorRT native SmolLM3 KV build requires PyTorch in the Model " + "Connect build environment to generate Hugging Face-exact RoPE " + "frequencies" + ) from exc + + with torch.no_grad(): + exponents = torch.arange( + 0, + rotary_ndims, + 2, + dtype=torch.int64, + device="cpu", + ).to(dtype=torch.float32) / rotary_ndims + inv_freq = 1.0 / (rope_theta ** exponents) + if rope_scaling: + rope_type = str( + rope_scaling.get("rope_type") + or rope_scaling.get("type") + or "" + ).lower() + if rope_type in ("", "default"): + rope_scaling = None + elif rope_type != "llama3": + raise ValueError( + "native active-position RoPE supports only llama3 scaling" + ) + if rope_scaling: + factor = float(rope_scaling["factor"]) + low = float(rope_scaling["low_freq_factor"]) + high = float(rope_scaling["high_freq_factor"]) + original = int( + rope_scaling["original_max_position_embeddings"] + ) + wavelength = 2 * math.pi / inv_freq + low_wavelength = original / low + high_wavelength = original / high + scaled = torch.where( + wavelength > low_wavelength, + inv_freq / factor, + inv_freq, + ) + smooth = ( + original / wavelength - low + ) / (high - low) + interpolated = ( + (1 - smooth) * scaled / factor + smooth * scaled + ) + medium = ( + ~(wavelength < high_wavelength) + * ~(wavelength > low_wavelength) + ) + inv_freq = torch.where(medium, interpolated, scaled) + return np.asarray( + inv_freq.detach().cpu().contiguous().numpy(), + dtype=np.float32, + ).copy() + + +def add_active_rope_cache( + network: trt.INetworkDefinition, + position_id: trt.ITensor, + inv_freq: np.ndarray, + output_dtype: trt.DataType, +) -> tuple[trt.ITensor, trt.ITensor]: + """Build ``[1, Sq, D/2]`` cos/sin tensors for active positions only.""" + inv_freq = np.asarray(inv_freq, dtype=np.float32) + if inv_freq.ndim != 1 or inv_freq.size == 0: + raise ValueError( + "active RoPE inverse frequencies must be a non-empty rank-1 array") + + pos_float = network.add_cast(position_id, trt.float32).get_output(0) + pos_col = network.add_shuffle(pos_float) + pos_col.reshape_dims = (-1, 1) + inv_freq_tensor = add_constant( + network, + (1, int(inv_freq.size)), + inv_freq.reshape(1, -1), + dtype=np.float32, + ) + angles = network.add_elementwise( + pos_col.get_output(0), + inv_freq_tensor, + trt.ElementWiseOperation.PROD, + ).get_output(0) + cos_2d = network.add_unary( + angles, trt.UnaryOperation.COS).get_output(0) + sin_2d = network.add_unary( + angles, trt.UnaryOperation.SIN).get_output(0) + + cos_3d = network.add_shuffle(cos_2d) + cos_3d.reshape_dims = (1, -1, int(inv_freq.size)) + sin_3d = network.add_shuffle(sin_2d) + sin_3d.reshape_dims = (1, -1, int(inv_freq.size)) + cos_cache = cos_3d.get_output(0) + sin_cache = sin_3d.get_output(0) + if cos_cache.dtype != output_dtype: + cos_cache = network.add_cast( + cos_cache, output_dtype).get_output(0) + sin_cache = network.add_cast( + sin_cache, output_dtype).get_output(0) + return cos_cache, sin_cache + + +def _yarn_inverse_frequencies( + inverse_frequencies: np.ndarray, + rotary_ndims: int, + rope_theta: float, + rope_scaling: Mapping[str, Any], +) -> tuple[np.ndarray, float]: + """YaRN-scaled inverse frequencies and the attention scale that goes with them. + + Mirrors Hugging Face's ``_compute_yarn_parameters``. Frequencies below the + ``beta_fast`` correction dimension keep their extrapolated value, those above + ``beta_slow`` are interpolated by ``factor``, and the band between the two is + ramped linearly. Upstream folds ``attention_factor`` into cos/sin rather than + into the attention scores, so it is returned here and applied to the table. + """ + factor = float(rope_scaling["factor"]) + original_context = float(rope_scaling["original_max_position_embeddings"]) + beta_fast = float(rope_scaling.get("beta_fast") or 32.0) + beta_slow = float(rope_scaling.get("beta_slow") or 1.0) + if ( + factor <= 0.0 + or original_context <= 0.0 + or beta_fast <= 0.0 + or beta_slow <= 0.0 + or beta_slow >= beta_fast + ): + raise ValueError(f"Invalid YaRN RoPE scaling: {dict(rope_scaling)}") + + def correction_dim(rotations: float) -> float: + return ( + rotary_ndims + * np.log(original_context / (rotations * 2.0 * np.pi)) + / (2.0 * np.log(rope_theta)) + ) + + half = rotary_ndims // 2 + low = max(int(np.floor(correction_dim(beta_fast))), 0) + high = min(int(np.ceil(correction_dim(beta_slow))), half - 1) + ramp = np.clip( + (np.arange(half, dtype=np.float64) - low) / max(high - low, 1), + 0.0, + 1.0, + ) + scaled = ( + inverse_frequencies / factor * ramp + + inverse_frequencies * (1.0 - ramp) + ) + attention_factor = rope_scaling.get("attention_factor") + if attention_factor is None: + attention_factor = 0.1 * np.log(factor) + 1.0 + return scaled, float(attention_factor) + + +def make_rope_table_half_dim( + max_cache_length: int, + head_dim: int, + rope_theta: float, + cosine: bool, + partial_rotary_factor: float = 1.0, + interleaved: bool = False, + rope_scaling: Mapping[str, Any] | None = None, +) -> np.ndarray: + """Build a RoPE cos/sin table of shape [max_cache_length, rotary_ndims // 2]. + + IRotaryEmbeddingLayer expects the cos/sin cache with only the *half* + rotary dimension (it internally handles both halves). This is different + from make_rope_table which produces [max_cache_length, hidden_size] by + repeating the per-head values across all heads. + + Args: + max_cache_length: Number of positions (rows in the table). + head_dim: Full head dimension (D). + rope_theta: Base frequency for inverse-frequency computation. + cosine: True → cos table, False → sin table. + partial_rotary_factor: Fraction of head dims that rotate (default 1.0). + interleaved: If True, adjacent-pair frequencies (CodeGen/GPT-J). + If False, half-split frequencies (LLaMA/Qwen). + rope_scaling: Optional Hugging Face RoPE scaling configuration. + + Returns: + Float32 array [max_cache_length, rotary_ndims // 2]. + """ + rotary_ndims = int(head_dim * partial_rotary_factor) + rotary_ndims = validate_native_rope_dim(rotary_ndims) + half = rotary_ndims // 2 + default = 1.0 if cosine else 0.0 + if max_cache_length <= 0 or rope_theta <= 0.0: + return np.full((max(max_cache_length, 1), max(half, 1)), + default, dtype=np.float32) + if not rope_scaling: + table = np.full( + (max_cache_length, half), + default, + dtype=np.float32, + ) + for pos in range(max_cache_length): + for dimension in range(half): + exponent = (2.0 * dimension) / rotary_ndims + inverse_frequency = rope_theta ** (-exponent) + angle = pos * inverse_frequency + table[pos, dimension] = ( + np.cos(angle) if cosine else np.sin(angle) + ) + return table + + exponents = np.arange(0, rotary_ndims, 2, dtype=np.float64) / rotary_ndims + inverse_frequencies = np.power(float(rope_theta), -exponents) + rope_type = str( + rope_scaling.get("rope_type") + or rope_scaling.get("type") + or "" + ).lower() + attention_scaling = 1.0 + if rope_type == "yarn": + inverse_frequencies, attention_scaling = _yarn_inverse_frequencies( + inverse_frequencies, rotary_ndims, float(rope_theta), rope_scaling + ) + angles = np.outer( + np.arange(max_cache_length, dtype=np.float64), + inverse_frequencies, + ) + values = np.cos(angles) if cosine else np.sin(angles) + return (values * attention_scaling).astype(np.float32) + + if rope_type != "llama3": + raise NotImplementedError( + f"Unsupported SmolLM3 RoPE scaling type: {rope_type or ''}" + ) + factor = float(rope_scaling["factor"]) + low_freq_factor = float(rope_scaling["low_freq_factor"]) + high_freq_factor = float(rope_scaling["high_freq_factor"]) + original_context = float( + rope_scaling["original_max_position_embeddings"] + ) + if ( + factor <= 0.0 + or low_freq_factor <= 0.0 + or high_freq_factor <= low_freq_factor + or original_context <= 0.0 + ): + raise ValueError(f"Invalid Llama-3 RoPE scaling: {dict(rope_scaling)}") + + wavelengths = 2.0 * np.pi / inverse_frequencies + low_freq_wavelength = original_context / low_freq_factor + high_freq_wavelength = original_context / high_freq_factor + scaled = np.where( + wavelengths > low_freq_wavelength, + inverse_frequencies / factor, + inverse_frequencies, + ) + smooth = ( + original_context / wavelengths - low_freq_factor + ) / (high_freq_factor - low_freq_factor) + interpolated = (1.0 - smooth) * scaled / factor + smooth * scaled + medium = ( + (wavelengths >= high_freq_wavelength) + & (wavelengths <= low_freq_wavelength) + ) + inverse_frequencies = np.where(medium, interpolated, scaled) + + angles = np.outer( + np.arange(max_cache_length, dtype=np.float64), + inverse_frequencies, + ) + values = np.cos(angles) if cosine else np.sin(angles) + return values.astype(np.float32) + + +def reshape_rows_to_heads_4d( + network: trt.INetworkDefinition, + x: trt.ITensor, + num_heads: int, + head_dim: int, + sequence_length: int | None = None, + tag: str | None = None, +) -> trt.ITensor: + """Reshape [S, H * D] rows into [1, H, S, D]. + + The transpose is required for S > 1 because each input row contains all + heads for one token. ``sequence_length=None`` means runtime-dynamic S. + """ + seq_dim = -1 if sequence_length is None else sequence_length + r1 = network.add_shuffle(x) + if tag: + r1.name = tag + "_s_h_d" + r1.reshape_dims = (seq_dim, num_heads, head_dim) + r1.second_transpose = trt.Permutation([1, 0, 2]) + + r2 = network.add_shuffle(r1.get_output(0)) + if tag: + r2.name = tag + "_1_h_s_d" + r2.reshape_dims = (1, num_heads, seq_dim, head_dim) + return r2.get_output(0) + + +def reshape_heads_4d_to_rows( + network: trt.INetworkDefinition, + x_4d: trt.ITensor, + attention_size: int, + sequence_length: int | None = None, + tag: str | None = None, +) -> trt.ITensor: + """Reshape [1, H, S, D] back to [S, H * D].""" + seq_dim = -1 if sequence_length is None else sequence_length + out = network.add_shuffle(x_4d) + if tag: + out.name = tag + "_s_h_d" + out.first_transpose = trt.Permutation([0, 2, 1, 3]) + out.reshape_dims = (seq_dim, attention_size) + return out.get_output(0) + + +def add_2d_mask_to_4d( + network: trt.INetworkDefinition, + mask_2d: trt.ITensor, +) -> trt.ITensor: + """Reshape additive attention mask [Sq, K] to [1, 1, Sq, K].""" + mask_shape = network.add_shape(mask_2d).get_output(0) + ones = add_constant( + network, (2,), np.array([1, 1], dtype=np.int64), dtype=np.int64) + target = network.add_concatenation([ones, mask_shape]) + target.axis = 0 + mask_4d = network.add_shuffle(mask_2d) + mask_4d.set_input(1, target.get_output(0)) + return mask_4d.get_output(0) + + +def add_alibi_mask_4d( + network: trt.INetworkDefinition, + mask_2d: trt.ITensor, + position_id: trt.ITensor, + alibi_slopes_tensor: trt.ITensor, + cache_position_indices: trt.ITensor, + num_heads: int, + target_dtype: trt.DataType | None = None, +) -> trt.ITensor: + """Build a per-head ALiBi additive mask for native IAttention. + + Args: + mask_2d: [Sq, K] additive mask. + position_id: [Sq] query positions. + alibi_slopes_tensor: [H, 1, 1] per-head slopes. + cache_position_indices: [cache_rows] key positions for cached rows. + target_dtype: Optional dtype for the returned mask. Defaults to + ``mask_2d.dtype``. + + Returns: + [1, H, Sq, K] additive mask containing both ``mask_2d`` and + ``slope[h] * (key_pos[k] - query_pos[q])``. + """ + pos_float = network.add_cast(position_id, trt.float32).get_output(0) + cache_positions = cache_position_indices + if cache_positions.dtype != trt.float32: + cache_positions = network.add_cast(cache_positions, trt.float32).get_output(0) + + key_pos = network.add_concatenation([cache_positions, pos_float]) + key_pos.axis = 0 + + mask_shape = network.add_shape(mask_2d).get_output(0) + one_const = add_constant( + network, (1,), np.array([1], dtype=np.int64), dtype=np.int64) + sq_size = network.add_slice(mask_shape, start=(0,), shape=(1,), stride=(1,)) + k_size = network.add_slice(mask_shape, start=(1,), shape=(1,), stride=(1,)) + sq_size_t = sq_size.get_output(0) + k_size_t = k_size.get_output(0) + + key_pos_shape = network.add_concatenation([one_const, k_size_t]) + key_pos_shape.axis = 0 + key_pos_2d = network.add_shuffle(key_pos.get_output(0)) + key_pos_2d.set_input(1, key_pos_shape.get_output(0)) + + query_pos_shape = network.add_concatenation([sq_size_t, one_const]) + query_pos_shape.axis = 0 + query_pos_2d = network.add_shuffle(pos_float) + query_pos_2d.set_input(1, query_pos_shape.get_output(0)) + + rel_pos = network.add_elementwise( + key_pos_2d.get_output(0), query_pos_2d.get_output(0), + trt.ElementWiseOperation.SUB) + + one_const2 = add_constant( + network, (1,), np.array([1], dtype=np.int64), dtype=np.int64) + rel_shape = network.add_concatenation([one_const, one_const2, sq_size_t, k_size_t]) + rel_shape.axis = 0 + rel_4d = network.add_shuffle(rel_pos.get_output(0)) + rel_4d.set_input(1, rel_shape.get_output(0)) + + slopes = alibi_slopes_tensor + if slopes.dtype != trt.float32: + slopes = network.add_cast(slopes, trt.float32).get_output(0) + slopes_4d = network.add_shuffle(slopes) + slopes_4d.reshape_dims = (1, num_heads, 1, 1) + + alibi_bias = network.add_elementwise( + slopes_4d.get_output(0), rel_4d.get_output(0), + trt.ElementWiseOperation.PROD) + alibi_bias_t = alibi_bias.get_output(0) + + mask_4d = add_2d_mask_to_4d(network, mask_2d) + out_dtype = target_dtype or mask_4d.dtype + if alibi_bias_t.dtype != out_dtype: + alibi_bias_t = network.add_cast(alibi_bias_t, out_dtype).get_output(0) + + combined = network.add_elementwise( + mask_4d, alibi_bias_t, trt.ElementWiseOperation.SUM) + return combined.get_output(0) + + +def add_apply_rope_native( + network: trt.INetworkDefinition, + inp: trt.ITensor, + num_heads: int, + head_dim: int, + cos_cache_2d: trt.ITensor, + sin_cache_2d: trt.ITensor, + position_id: trt.ITensor | None, + rotary_embedding_dim: int, + interleaved: bool = False, + sequence_length: int | None = 1, +) -> trt.ITensor: + """Apply RoPE via TRT native IRotaryEmbeddingLayer. + + Handles both single-token decoder steps and dynamic-Sq prefill/decode + graphs without a manual rotate-half matmul chain. + + Shape contract with indexed, full-capacity caches: + input: [1, num_heads, Sq, head_dim] (reshaped internally) + cos_cache_2d: [max_S, rotary_embedding_dim // 2] (2-D constant) + sin_cache_2d: [max_S, rotary_embedding_dim // 2] (2-D constant) + position_id: [Sq] int32, reshaped to [1, Sq] internally + + Shape contract with active-position caches: + input: [1, num_heads, Sq, head_dim] + cos_cache_2d: [1, Sq, rotary_embedding_dim // 2] + sin_cache_2d: [1, Sq, rotary_embedding_dim // 2] + position_id: None; the caches already correspond to active positions + interleaved: False → rotate-half (LLaMA/Qwen) + True → adjacent-pair (CodeGen/GPT-J) + + Args: + inp: [Sq, num_heads * head_dim]. + num_heads: Number of attention heads. + head_dim: Per-head dimension. + cos_cache_2d: Pre-built 2-D cos table constant. + sin_cache_2d: Pre-built 2-D sin table constant. + position_id: Runtime position indices, shape [Sq] int32, or + ``None`` for active-position rank-3 caches. + rotary_embedding_dim: Number of head dims that participate in RoPE. + interleaved: Frequency layout (see above). + sequence_length: Static Sq, or None for runtime-dynamic Sq. + + Returns: + [Sq, num_heads * head_dim] with RoPE applied. + """ + rotary_embedding_dim = validate_native_rope_dim(rotary_embedding_dim) + attention_size = num_heads * head_dim + + inp_4d = reshape_rows_to_heads_4d( + network, inp, num_heads, head_dim, sequence_length) + + rope = network.add_rotary_embedding( + inp_4d, + cos_cache_2d, + sin_cache_2d, + interleaved, + rotary_embedding_dim, + ) + if position_id is not None: + # Reshape position_id [Sq] -> [1, Sq] (batch=1). + seq_dim = -1 if sequence_length is None else sequence_length + pos_2d = network.add_shuffle(position_id) + pos_2d.reshape_dims = (1, seq_dim) + rope.set_input(3, pos_2d.get_output(0)) + + return reshape_heads_4d_to_rows( + network, rope.get_output(0), attention_size, sequence_length) + + +def add_attention_core( + network: trt.INetworkDefinition, + q_4d: trt.ITensor, + k_4d: trt.ITensor, + v_4d: trt.ITensor, + causal: bool = False, + mask: trt.ITensor | None = None, + scale: float | None = None, + fp32_accumulation: bool = False, +) -> trt.ITensor: + """Scaled dot-product attention via TRT native IAttention layer. + + Replaces the manual Q@K^T → scale → softmax → @V chain. TRT 10 fuses + this into a single kernel when a compatible implementation is available; + decomposable=True ensures a correct fallback to primitives otherwise. + + NOTE: TRT IAttention computes raw BMM1 = Q @ K^T without any built-in + 1/sqrt(D) scaling. We pre-scale Q by 1/sqrt(D) so that the fused kernel + computes the standard scaled dot-product attention formula. + + Args: + q_4d: Query [B, H, q_seq, D]. + k_4d: Key [B, H, kv_seq, D]. + v_4d: Value [B, H, kv_seq, D]. + causal: Apply causal (autoregressive) mask. Mutually exclusive + with ``mask``. + mask: Optional additive float mask [B, H, q_seq, kv_seq] added + to scaled logits before softmax. Cannot be used with + causal=True. + scale: Optional Q pre-scale factor. Defaults to 1/sqrt(D). + fp32_accumulation: + Cast Q/K/V to FP32 before IAttention, then cast the context + back to the original Q dtype. TRT may still select a + Half-input fused MHA tactic after optimizing the casts, while + keeping the IAttention accumulation/output boundary in FP32. + + Returns: + Context tensor [B, H, q_seq, D]. + """ + output_dtype = q_4d.dtype + if fp32_accumulation and output_dtype != trt.float32: + q_4d = network.add_cast(q_4d, trt.float32).get_output(0) + k_4d = network.add_cast(k_4d, trt.float32).get_output(0) + v_4d = network.add_cast(v_4d, trt.float32).get_output(0) + if mask is not None and mask.dtype != trt.float32: + mask = network.add_cast(mask, trt.float32).get_output(0) + + # Pre-scale Q: TRT IAttention does not apply score scaling itself. + # Match the scale constant's dtype to Q's dtype: in strongly-typed networks + # a FP32 constant mixed with a FP16/BF16 Q causes add_elementwise to emit + # a type-mismatch error and produce a tensor with corrupted dimensions, + # which makes add_attention return None. + if scale is None: + head_dim = q_4d.shape[-1] + scale = float(1.0 / np.sqrt(head_dim)) if head_dim > 0 else 1.0 + # Use FP16 weights directly for FP16; BF16 has no numpy native type so + # create as FP32 and cast; FP32 falls through to the default. + scale_np_dtype = np.float16 if q_4d.dtype == trt.float16 else np.float32 + scale_t = add_constant(network, (1, 1, 1, 1), np.array([[[[scale]]]]), dtype=scale_np_dtype) + if q_4d.dtype == trt.bfloat16: + scale_t = network.add_cast(scale_t, trt.bfloat16).get_output(0) + q_scaled = network.add_elementwise(q_4d, scale_t, trt.ElementWiseOperation.PROD) + + attn = network.add_attention( + q_scaled.get_output(0), k_4d, v_4d, + trt.AttentionNormalizationOp.SOFTMAX, + causal, + ) + # Allow TRT to decompose into primitive ops when no fused kernel is + # available (e.g. unsupported head-dim or dtype). This guarantees + # correctness on any configuration at the cost of potential performance. + attn.decomposable = True + if mask is not None and not causal: + attn.mask = mask + return _cast_back_to_trt_dtype(network, attn.get_output(0), output_dtype) + + +def add_native_kv_cache_attention_from_rows( + network: trt.INetworkDefinition, + q: trt.ITensor, + k_update: trt.ITensor, + v_update: trt.ITensor, + cache_k: trt.ITensor, + cache_v: trt.ITensor, + cache_write_indices: trt.ITensor, + attention_masks: NativeKvMasks, + *, + num_heads: int, + num_kv_heads: int, + head_dim: int, + q_seq: int | None, + scale: float | None = None, + tag: str | None = None, +) -> dict[str, trt.ITensor]: + """Update a user-owned KV cache and attend over its active prefix. + + ``IKVCacheUpdateLayer`` writes this step's K/V rows into the static, + full-capacity cache in place. Attention uses a shared explicit + active-prefix causal mask and a primitive grouped-query graph, avoiding + correctness and kernel-coverage dependence on + ``IAttention.key_value_lengths``. + + Cache inputs and outputs have shape ``[1, Hkv, capacity, D]``. The runtime + must bind each output to the same device address as its corresponding + input, as required by TensorRT's KV-cache aliasing contract. + """ + if not hasattr(network, "add_kv_cache_update"): + raise RuntimeError( + "SmolLM3 native KV cache requires TensorRT add_kv_cache_update support" + ) + + k_update_4d = reshape_rows_to_heads_4d( + network, + k_update, + num_kv_heads, + head_dim, + sequence_length=q_seq, + tag=None if tag is None else tag + ".k_update", + ) + v_update_4d = reshape_rows_to_heads_4d( + network, + v_update, + num_kv_heads, + head_dim, + sequence_length=q_seq, + tag=None if tag is None else tag + ".v_update", + ) + + update_k = network.add_kv_cache_update( + cache_k, + k_update_4d, + cache_write_indices, + trt.KVCacheMode.LINEAR, + ) + update_v = network.add_kv_cache_update( + cache_v, + v_update_4d, + cache_write_indices, + trt.KVCacheMode.LINEAR, + ) + if update_k is None or update_v is None: + raise RuntimeError("TensorRT failed to create SmolLM3 KV-cache update layers") + if tag: + update_k.name = tag + ".cache_k_update" + update_v.name = tag + ".cache_v_update" + updated_k = update_k.get_output(0) + updated_v = update_v.get_output(0) + + q_4d = reshape_rows_to_heads_4d( + network, + q, + num_heads, + head_dim, + sequence_length=q_seq, + tag=None if tag is None else tag + ".q", + ) + if q_4d.dtype != trt.bfloat16: + raise ValueError("SmolLM3 native KV attention requires BF16 queries") + context_4d = add_explicit_masked_grouped_query_attention( + network, + q_4d, + updated_k, + updated_v, + attention_masks, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + scale=scale, + tag=tag, + ) + context = reshape_heads_4d_to_rows( + network, + context_4d, + num_heads * head_dim, + sequence_length=q_seq, + tag=None if tag is None else tag + ".ctx", + ) + return { + "context": context, + "present_k": updated_k, + "present_v": updated_v, + } + + +def _scalar_constant_for_trt_dtype( + network: trt.INetworkDefinition, + shape: tuple[int, ...], + value: float, + dtype: trt.DataType, +) -> trt.ITensor: + np_dtype = np.float16 if dtype == trt.float16 else np.float32 + const = add_constant( + network, shape, np.full(shape, value, dtype=np_dtype), + dtype=np_dtype) + if dtype == trt.bfloat16: + const = network.add_cast(const, trt.bfloat16).get_output(0) + return const + + +def add_tanh_softcap( + network: trt.INetworkDefinition, + tensor: trt.ITensor, + cap: float, + *, + scalar_shape: tuple[int, ...], +) -> trt.ITensor: + """Apply ``tanh(tensor / cap) * cap`` using scalar broadcasting.""" + cap_t = _scalar_constant_for_trt_dtype( + network, scalar_shape, float(cap), tensor.dtype) + scaled = network.add_elementwise( + tensor, cap_t, trt.ElementWiseOperation.DIV).get_output(0) + capped = network.add_activation( + scaled, trt.ActivationType.TANH).get_output(0) + return network.add_elementwise( + capped, cap_t, trt.ElementWiseOperation.PROD).get_output(0) + + +def _repeat_kv_heads_4d( + network: trt.INetworkDefinition, + x_4d: trt.ITensor, + *, + num_heads: int, + num_kv_heads: int, + head_dim: int, +) -> trt.ITensor: + if num_kv_heads == num_heads: + return x_4d + if num_kv_heads <= 0 or num_heads % num_kv_heads != 0: + raise ValueError( + f"num_heads={num_heads} must be divisible by " + f"num_kv_heads={num_kv_heads}") + + repeat = num_heads // num_kv_heads + if num_kv_heads == 1: + concat = network.add_concatenation([x_4d] * repeat) + concat.axis = 1 + return concat.get_output(0) + + x_shape = network.add_shape(x_4d).get_output(0) + one = add_constant( + network, (1,), np.array([1], dtype=np.int64), dtype=np.int64) + seq = network.add_slice(x_shape, start=(2,), shape=(1,), stride=(1,)) + dim = add_constant( + network, (1,), np.array([head_dim], dtype=np.int64), dtype=np.int64) + slice_shape = network.add_concatenation([one, one, seq.get_output(0), dim]) + slice_shape.axis = 0 + + repeated = [] + for head_idx in range(num_kv_heads): + head_slice = network.add_slice( + x_4d, start=(0, head_idx, 0, 0), + shape=(1, 1, 1, head_dim), stride=(1, 1, 1, 1)) + head_slice.set_input(2, slice_shape.get_output(0)) + repeated.extend([head_slice.get_output(0)] * repeat) + + concat = network.add_concatenation(repeated) + concat.axis = 1 + return concat.get_output(0) + + +def _add_attention_core_with_logit_softcap( + network: trt.INetworkDefinition, + q_4d: trt.ITensor, + k_4d: trt.ITensor, + v_4d: trt.ITensor, + *, + num_heads: int, + num_kv_heads: int, + head_dim: int, + mask: trt.ITensor | None, + scale: float, + logit_softcap: float, +) -> trt.ITensor: + output_dtype = q_4d.dtype + k_4d = _repeat_kv_heads_4d( + network, k_4d, num_heads=num_heads, num_kv_heads=num_kv_heads, + head_dim=head_dim) + v_4d = _repeat_kv_heads_4d( + network, v_4d, num_heads=num_heads, num_kv_heads=num_kv_heads, + head_dim=head_dim) + + score_q = q_4d + score_k = k_4d + score_mask = mask + if output_dtype != trt.float32: + score_q = network.add_cast(score_q, trt.float32).get_output(0) + score_k = network.add_cast(score_k, trt.float32).get_output(0) + if score_mask is not None and score_mask.dtype != trt.float32: + score_mask = network.add_cast(score_mask, trt.float32).get_output(0) + + scale_t = _scalar_constant_for_trt_dtype( + network, (1, 1, 1, 1), scale, score_q.dtype) + scores = network.add_matrix_multiply( + score_q, trt.MatrixOperation.NONE, + score_k, trt.MatrixOperation.TRANSPOSE).get_output(0) + scores = network.add_elementwise( + scores, scale_t, trt.ElementWiseOperation.PROD).get_output(0) + + scores = add_tanh_softcap( + network, scores, logit_softcap, scalar_shape=(1, 1, 1, 1)) + + if score_mask is not None: + scores = network.add_elementwise( + scores, score_mask, trt.ElementWiseOperation.SUM).get_output(0) + + probs = network.add_softmax(scores) + probs.axes = 1 << 3 + probs_t = probs.get_output(0) + if probs_t.dtype != output_dtype: + probs_t = network.add_cast(probs_t, output_dtype).get_output(0) + + context = network.add_matrix_multiply( + probs_t, trt.MatrixOperation.NONE, + v_4d, trt.MatrixOperation.NONE).get_output(0) + return _cast_back_to_trt_dtype(network, context, output_dtype) + + +def add_attention_from_rows( + network: trt.INetworkDefinition, + q: trt.ITensor, + k: trt.ITensor, + v: trt.ITensor, + *, + num_heads: int, + head_dim: int, + num_kv_heads: int | None = None, + q_seq: int | None, + kv_seq: int | None, + causal: bool = False, + mask: trt.ITensor | None = None, + scale: float | None = None, + logit_softcap: float | None = None, + fp32_accumulation: bool = False, + tag: str | None = None, +) -> trt.ITensor: + """Native IAttention for row-major [S, H * D] Q/K/V tensors. + + ``num_kv_heads`` can be smaller than ``num_heads`` for GQA/MQA. TRT + native IAttention supports this directly, so callers should not expand K/V + heads unless the model semantics require per-query-head K/V values. + """ + attention_size = num_heads * head_dim + kv_heads = num_heads if num_kv_heads is None else num_kv_heads + q_4d = reshape_rows_to_heads_4d( + network, q, num_heads, head_dim, sequence_length=q_seq, + tag=None if tag is None else tag + ".q") + k_4d = reshape_rows_to_heads_4d( + network, k, kv_heads, head_dim, sequence_length=kv_seq, + tag=None if tag is None else tag + ".k") + v_4d = reshape_rows_to_heads_4d( + network, v, kv_heads, head_dim, sequence_length=kv_seq, + tag=None if tag is None else tag + ".v") + if scale is None: + scale = float(1.0 / np.sqrt(head_dim)) if head_dim > 0 else 1.0 + if logit_softcap is not None and float(logit_softcap) > 0.0: + if causal: + raise NotImplementedError( + "logit_softcap attention requires an explicit additive mask") + ctx_4d = _add_attention_core_with_logit_softcap( + network, q_4d, k_4d, v_4d, + num_heads=num_heads, num_kv_heads=kv_heads, head_dim=head_dim, + mask=mask, scale=scale, logit_softcap=float(logit_softcap)) + else: + ctx_4d = add_attention_core( + network, q_4d, k_4d, v_4d, causal=causal, mask=mask, scale=scale, + fp32_accumulation=fp32_accumulation) + return reshape_heads_4d_to_rows( + network, ctx_4d, attention_size, sequence_length=q_seq, + tag=None if tag is None else tag + ".ctx") + + +# Backward-compatible name used by existing tests and call sites. +_add_attention_core = add_attention_core + + +def add_decoder_attention_ffi( + network: trt.INetworkDefinition, + q: trt.ITensor, + all_k: trt.ITensor, + all_v: trt.ITensor, + *, + kernel_name: str, + num_heads: int, + head_dim: int, + attention_window: int, +) -> trt.ITensor: + """Decoder attention via TVM-FFI kernel (FlashInfer, CuTe, etc). + + The kernel must be registered as a TVM-FFI global before engine build. + + Inputs: + q: [1, attention_size] + all_k, all_v: [attention_window, attention_size] + Returns: + context: [1, attention_size] + """ + attention_size = num_heads * head_dim + + q_2d = network.add_shuffle(q) + q_2d.reshape_dims = (num_heads, head_dim) + k_3d = network.add_shuffle(all_k) + k_3d.reshape_dims = (attention_window, num_heads, head_dim) + v_3d = network.add_shuffle(all_v) + v_3d.reshape_dims = (attention_window, num_heads, head_dim) + + scale_val = 1.0 / (head_dim ** 0.5) + ffi_outputs = add_tvm_ffi_kernel( + network, + kernel_name=kernel_name, + inputs=[q_2d.get_output(0), k_3d.get_output(0), + v_3d.get_output(0)], + output_specs=[{"dims": [num_heads, head_dim], "dtype": "float16"}], + workspace_bytes=32 * 1024 * 1024, # 32MB for FlashInfer tmp + extra_args=[ + {"type": "none"}, # maybe_lse + {"type": "int", "value": 0}, # kv_layout_code (NHD) + {"type": "int", "value": -1}, # window_left + {"type": "none"}, # alibi_slopes + {"type": "float", "value": 0.0}, # logits_soft_cap + {"type": "float", "value": scale_val}, # sm_scale + {"type": "float", "value": 1.0}, # rope_rcp_scale + {"type": "float", "value": 0.0001}, # rope_rcp_theta + ], + ) + context_flat = network.add_shuffle(ffi_outputs[0]) + context_flat.reshape_dims = (1, attention_size) + return context_flat.get_output(0) + + +# --------------------------------------------------------------------------- +# TVM-FFI kernel bridge +# --------------------------------------------------------------------------- + +def add_tvm_ffi_kernel( + network: trt.INetworkDefinition, + kernel_name: str, + inputs: list[trt.ITensor], + output_specs: list[dict], + workspace_bytes: int = 0, + extra_args: list[dict] | None = None, +) -> list[trt.ITensor]: + """Add a TVM-FFI kernel call as a TRT plugin layer. + + Args: + network: TRT network being built. + kernel_name: TVM-FFI global function name (e.g. "my_ns.my_kernel"). + inputs: List of input ITensor objects. + output_specs: List of dicts, one per output. Each dict has: + - "dims": "same_as_input_N" or list of ints for fixed shape + - "dtype": "float32" or "float16" (default "float32") + workspace_bytes: Extra workspace bytes for the kernel (default 0). + extra_args: Optional list of extra scalar/pointer args to pass after + tensors. Each dict has "type" ("none"|"int"|"float"|"ptr") and + optional "value". + + Returns: + List of output ITensor objects. + """ + import json + + registry = trt.get_plugin_registry() + creator = registry.get_plugin_creator("TvmFfiKernel", "1", "") + if creator is None: + raise RuntimeError( + "TvmFfiKernel plugin not found in TRT registry. " + "Ensure the C++ plugin is compiled with TRTMC_HAS_TVM_FFI=1." + ) + + spec_dict = { + "num_inputs": len(inputs), + "num_outputs": len(output_specs), + "outputs": output_specs, + "workspace_bytes": workspace_bytes, + } + if extra_args: + spec_dict["extra_args"] = extra_args + shape_spec = json.dumps(spec_dict) + + fields = [ + trt.PluginField("kernel_name", kernel_name.encode("utf-8"), + trt.PluginFieldType.CHAR), + trt.PluginField("shape_spec", shape_spec.encode("utf-8"), + trt.PluginFieldType.CHAR), + ] + fc = trt.PluginFieldCollection(fields) + plugin = creator.create_plugin("tvm_ffi_kernel", fc) + + layer = network.add_plugin_v2(inputs, plugin) + return [layer.get_output(i) for i in range(layer.num_outputs)] diff --git a/python/tensorrt_model_connect/families/smollm3/native_kv_contract.py b/python/tensorrt_model_connect/families/smollm3/native_kv_contract.py new file mode 100644 index 0000000000..7fca8f2246 --- /dev/null +++ b/python/tensorrt_model_connect/families/smollm3/native_kv_contract.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Mapped-weight contract for SmolLM3's TensorRT native KV graph.""" + +from __future__ import annotations + +import re +from collections.abc import Mapping + +from .build_routing import resolved_head_dim + +_LAYER_KEY = re.compile(r"^layer\.(\d+)\.") +_BIAS_SUFFIXES = ( + "input_norm_beta", + "q_bias", + "k_bias", + "v_bias", + "o_bias", + "post_attn_norm_beta", + "gate_bias", + "up_bias", + "down_bias", +) + + +def _shape(value: object, name: str) -> tuple[int, ...]: + try: + return tuple(int(dim) for dim in value.shape) + except (AttributeError, TypeError, ValueError, OverflowError) as exc: + raise ValueError( + f"native SmolLM3 weight {name} has an invalid shape" + ) from exc + + +def _require( + weights: Mapping[str, object], + name: str, + expected: tuple[int, ...], +) -> None: + if name not in weights: + raise ValueError(f"missing native SmolLM3 weight {name}") + actual = _shape(weights[name], name) + if actual != expected: + raise ValueError( + f"native SmolLM3 weight {name} must have shape " + f"{expected}, got {actual}" + ) + + +def validate_native_kv_weights( + config: object, + weights: Mapping[str, object], +) -> None: + """Validate mapped tensors once, before TensorRT graph construction.""" + + if not isinstance(weights, Mapping): + raise ValueError("native SmolLM3 weights must be a mapping") + + hidden = int(getattr(config, "hidden_size")) + vocab = int(getattr(config, "vocab_size")) + mlp = int(getattr(config, "intermediate_size")) + layers = int(getattr(config, "num_hidden_layers")) + heads = int(getattr(config, "num_attention_heads")) + kv_heads = int(getattr(config, "num_key_value_heads")) + head_dim = resolved_head_dim(config) + attention = heads * head_dim + kv_attention = kv_heads * head_dim + + layer_indices: set[int] = set() + malformed: list[str] = [] + for name in weights: + if not isinstance(name, str) or not name.startswith("layer."): + continue + match = _LAYER_KEY.match(name) + if match is None: + malformed.append(name) + else: + layer_indices.add(int(match.group(1))) + if malformed or layer_indices != set(range(layers)): + raise ValueError( + "native SmolLM3 weights require continuous layer indices; " + f"malformed={sorted(malformed)}, found={sorted(layer_indices)}" + ) + + for name, expected in ( + ("_attention_size", attention), + ("_kv_attention_size", kv_attention), + ("_mlp_size", mlp), + ): + if name in weights and int(weights[name]) != expected: + raise ValueError( + f"native SmolLM3 metadata {name} must be {expected}" + ) + + _require(weights, "embedding", (vocab, hidden)) + _require(weights, "final_norm", (hidden,)) + _require(weights, "w_out", (hidden, vocab)) + forbidden = [ + name + for name in ("final_norm_beta", "lm_head_bias") + if name in weights + ] + + for layer in range(layers): + prefix = f"layer.{layer}" + for suffix, expected in ( + ("input_norm", (hidden,)), + ("w_q", (hidden, attention)), + ("w_k", (hidden, kv_attention)), + ("w_v", (hidden, kv_attention)), + ("w_o", (attention, hidden)), + ("post_attn_norm", (hidden,)), + ("w_gate", (hidden, mlp)), + ("w_up", (hidden, mlp)), + ("w_down", (mlp, hidden)), + ): + _require(weights, f"{prefix}.{suffix}", expected) + for suffix, expected in ( + ("q_norm", (attention,)), + ("k_norm", (kv_attention,)), + ): + name = f"{prefix}.{suffix}" + if name in weights: + _require(weights, name, expected) + forbidden.extend( + f"{prefix}.{suffix}" + for suffix in _BIAS_SUFFIXES + if f"{prefix}.{suffix}" in weights + ) + + if forbidden: + raise ValueError( + "native dense SmolLM3 does not support bias weights: " + + ", ".join(sorted(forbidden)) + ) diff --git a/python/tensorrt_model_connect/families/smollm3/plugin.py b/python/tensorrt_model_connect/families/smollm3/plugin.py new file mode 100644 index 0000000000..59bf49e0b9 --- /dev/null +++ b/python/tensorrt_model_connect/families/smollm3/plugin.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""SmolLM3 family plugin.""" + +from __future__ import annotations + +from .config import ModelConfig +from .checkpoint_mapper import WeightDict, load_standard_weights +from .build_routing import ( + native_kv_architecture_capability, + native_kv_build_capability, +) +from .native_kv_contract import validate_native_kv_weights +from .dual_profile_decoder_builder import build_dual_profile_decoder_engine +from .standard_decoder_builder import build_standard_decoder_engine + + +class SmolLM3Plugin: + name = "smollm3" + runtime_strategy = "smollm3_decoder_kv_cache" + runtime_capabilities = {"decoder_kv"} + + def matches(self, model_type: str) -> bool: + return model_type.lower() == "smollm3" + + def default_build_precision(self, config: ModelConfig) -> str: + capability = native_kv_architecture_capability(config) + return "bf16" if capability.eligible else "fp32" + + def default_max_cache_length(self, config: ModelConfig) -> int: + """Use the model's complete context for native SmolLM3.""" + capability = native_kv_architecture_capability(config) + return int(config.max_position_embeddings) if capability.eligible else 256 + + def supports_split_decoder_roles(self, config: ModelConfig) -> bool: + return not bool(config.raw.get("_fp32_layers")) + + def load_weights( + self, model_dir: str, config: ModelConfig, + *, precision: str = "fp32", + ) -> WeightDict: + return load_standard_weights( + model_dir, + config, + precision=precision, + fp32_layers=tuple(config.raw.get("_fp32_layers", ())), + ) + + def build_engine( + self, config: ModelConfig, weights: WeightDict, + max_cache_length: int, *, precision: str = "fp32", + quant_ctx=None, verbose: bool = False, debug_layer_outputs: bool = False, + ) -> bytes: + capability = native_kv_build_capability( + config, + precision=precision, + max_cache_length=max_cache_length, + quantized=quant_ctx is not None, + debug_layer_outputs=debug_layer_outputs, + ) + if capability.eligible: + validate_native_kv_weights(config, weights) + config.raw["_decoder_engine_layout_supported"] = True + config.raw["_native_kv_cache_metadata"] = { + "native_kv_contract_version": 1, + "native_kv_cache": True, + } + role = str( + config.raw.get("_decoder_engine_role", "") + ) + if role not in ("prefill", "decode"): + raise ValueError( + "native SmolLM3 requires explicit split engine role " + f"'prefill' or 'decode', got {role!r}" + ) + return build_dual_profile_decoder_engine( + config, + weights, + max_cache_length, + precision="bf16", + quant_ctx=None, + verbose=verbose, + profile_mode=role, + native_kv_cache=True, + ) + + config.raw.pop("_native_kv_cache_metadata", None) + return build_standard_decoder_engine( + config, weights, max_cache_length, precision=precision, + quant_ctx=quant_ctx, verbose=verbose, + debug_layer_outputs=debug_layer_outputs) + + def get_bundle_config_overrides( + self, config: ModelConfig, + ) -> dict | None: + """Mark bundles that use the native KV runtime contract.""" + metadata = config.raw.get("_native_kv_cache_metadata") + return dict(metadata) if isinstance(metadata, dict) else None + + +plugin = SmolLM3Plugin() diff --git a/python/tensorrt_model_connect/families/smollm3/standard_decoder_builder.py b/python/tensorrt_model_connect/families/smollm3/standard_decoder_builder.py new file mode 100644 index 0000000000..b26e7047bb --- /dev/null +++ b/python/tensorrt_model_connect/families/smollm3/standard_decoder_builder.py @@ -0,0 +1,694 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Standard decoder engine builder (parameterized). + +Builds a TensorRT engine for a decoder-only transformer. Supports multiple +norm, MLP, and position embedding strategies via parameters: + + norm_type: "rmsnorm" | "layernorm" + mlp_type: "swiglu" | "gelu_fc" + position_type: "rope" | "learned" | "alibi" + activation: "silu" | "gelu_new" | "gelu" | "relu" | "relu2" (used by gelu_fc MLP) + +Tensor names MUST match what the C++ runtime expects: + Inputs: token_id, position_id, attention_mask, cache_k_0..N, cache_v_0..N + Outputs: logits, present_k_0..N, present_v_0..N +""" + +from __future__ import annotations + +import sys +from typing import TYPE_CHECKING + +import numpy as np +from tensorrt_model_connect import trt_compat + +from . import graph_ops +from . import graph_blocks +from .config import ModelConfig, resolve_rope_layer_schedule +from .dual_profile_decoder_builder import build_dual_profile_decoder_engine +from .utils import const_in_work_dtype + +trt = trt_compat.get_trt() + +if TYPE_CHECKING: + from .checkpoint_mapper import WeightDict + from ...quantization.context import QuantContext + + +def _mark_debug_output( + network: trt.INetworkDefinition, + tensor: trt.ITensor, + name: str, +) -> None: + """Mark a tensor as a network output for debug inspection.""" + # Use an identity layer to avoid aliasing issues with existing outputs. + cast = network.add_cast(tensor, trt.float32) + out = cast.get_output(0) + out.name = name + network.mark_output(out) + + +def build_standard_decoder_engine( + config: ModelConfig, + weights: WeightDict, + max_cache_length: int, + *, + precision: str = "fp32", + quant_ctx: QuantContext | None = None, + norm_type: str = "rmsnorm", + mlp_type: str = "swiglu", + position_type: str = "rope", + activation: str = "silu", + partial_rotary_factor: float = 1.0, + interleaved_rope: bool = False, + parallel_residual: bool = False, + scale_attn_weights: bool = True, + alibi_bias_scale: float = 1.0, + embed_input: bool = False, + verbose: bool = False, + debug_layer_outputs: bool = False, + hidden_state_output: bool = False, +) -> bytes: + """Build a TRT engine plan (serialized bytes) for a standard decoder. + + Args: + config: Model architecture from config.json. + weights: Loaded weight dict from checkpoint_mapper. + max_cache_length: KV cache length (engine is compiled for this value). + precision: Compute precision ("fp32", "fp16", or "bf16"). + norm_type: "rmsnorm" or "layernorm". + mlp_type: "swiglu" (3 projections: gate/up/down) or + "gelu_fc" (2 projections: fc1/fc2 with activation). + position_type: "rope" (rotary), "learned" (absolute position embeddings), + or "alibi" (attention with linear biases, no position embeddings). + activation: Activation function for gelu_fc MLP ("gelu_new", "gelu", "relu", "relu2"). + partial_rotary_factor: Fraction of head dims that get RoPE (default 1.0). + interleaved_rope: If True, use interleaved RoPE (CodeGen/GPT-J) where + adjacent dims (d, d+1) share frequencies. Default False uses + rotated-half (LLaMA/Qwen) where (d, d+half) share frequencies. + scale_attn_weights: Whether to scale attention scores by 1/sqrt(head_dim). + Most models use this (True, default). GPT-Neo does NOT scale (False). + alibi_bias_scale: Extra scale applied to ALiBi slopes before building + the additive attention mask. BLOOM leaves ALiBi unscaled while + Falcon scales ALiBi with the same factor as QK scores. + embed_input: If True, add input_embed [1, hidden] and use_input_embed [1] + engine inputs. When use_input_embed==1, the decoder uses input_embed + directly instead of the embedding lookup. Used for VL models where + the vision encoder provides fused embeddings during prefill. + verbose: Print TRT builder logs. + debug_layer_outputs: If True, mark per-layer hidden states as network + outputs for diff testing. + + Returns: + Serialized engine plan bytes. + """ + import os as _os + # Mark the graph as honoring the internal decoder role contract. This is + # embedded in the mutable config for family helpers that need to branch on + # the active engine layout while building. + config.raw["_decoder_engine_layout_supported"] = True + decoder_engine_role = str(config.raw.get("_decoder_engine_role", "dual_profile")) + + # Dispatch to the dynamic-Sq builder for dual-profile and split-prefill + # engines. Quantized builds (``quant_ctx``) thread Q/DQ insertion through + # every projection matmul via + # ``QuantContext.maybe_quantized_matmul``, so they share the dispatch. + # + # The legacy single-profile graph below stays in place for paths the + # dual-profile builder does not yet cover: + # + # - embed_input=True (VL prefill replacement, Bark sub-engines) + # - debug_layer_outputs=True (per-layer hidden-state dumps) + # - hidden_state_output=True (speech / Bark hidden output) + # + # ``TRTMC_NO_DUAL_PROFILE=1`` is an internal escape hatch (perf A/B, + # bisects against the legacy graph). It is *not* intended as a + # supported user-facing flag. + requested_fp32_layers = tuple(config.raw.get("_fp32_layers", ())) + dynamic_kv_cache = bool(config.raw.get("dynamic_kv_cache", False)) + if dynamic_kv_cache and position_type == "alibi": + raise ValueError("dynamic_kv_cache is not supported for ALiBi decoder builds") + dynamic_kv_profile_rows = ( + config.raw.get("_dynamic_kv_profile_rows") if dynamic_kv_cache else None + ) + if dynamic_kv_cache and not dynamic_kv_profile_rows: + dynamic_kv_profile_rows = [max_cache_length] + _dual_profile_disabled_for = ( + embed_input + or debug_layer_outputs + or hidden_state_output + or bool(requested_fp32_layers) + or _os.environ.get("TRTMC_NO_DUAL_PROFILE") == "1" + ) + if decoder_engine_role == "prefill" and _dual_profile_disabled_for: + raise NotImplementedError( + "split prefill engine is not supported for this standard decoder " + "configuration") + if not _dual_profile_disabled_for and decoder_engine_role in ("dual_profile", "prefill"): + return build_dual_profile_decoder_engine( + config, weights, max_cache_length, + precision=precision, + quant_ctx=quant_ctx, + norm_type=norm_type, + mlp_type=mlp_type, + position_type=position_type, + activation=activation, + partial_rotary_factor=partial_rotary_factor, + interleaved_rope=interleaved_rope, + parallel_residual=parallel_residual, + scale_attn_weights=scale_attn_weights, + alibi_bias_scale=alibi_bias_scale, + verbose=verbose, + dynamic_kv_profile_rows=dynamic_kv_profile_rows, + profile_mode=("prefill" if decoder_engine_role == "prefill" else "dual_profile"), + ) + + attention_size: int = weights.get("_attention_size", config.attention_size) + mlp_size: int = weights.get("_mlp_size", config.intermediate_size) + hidden = config.hidden_size + vocab = config.vocab_size + num_layers = config.num_hidden_layers + num_heads = config.num_attention_heads + num_kv_heads = config.num_key_value_heads + fp32_layers = frozenset(int(layer) for layer in requested_fp32_layers) + invalid_fp32_layers = sorted( + layer for layer in fp32_layers if layer < 0 or layer >= num_layers) + if invalid_fp32_layers: + raise ValueError( + f"fp32_layers contains out-of-range indices: {invalid_fp32_layers}") + if precision == "fp32": + fp32_layers = frozenset() + if fp32_layers and quant_ctx is not None: + raise ValueError("fp32_layers is not supported with quantized builds") + head_dim = attention_size // num_heads + kv_attention_size = graph_blocks.infer_kv_attention_size( + weights, num_kv_heads=num_kv_heads, head_dim=head_dim) + attention_window = max_cache_length + 1 + dynamic_kv_opt_rows = int(config.raw.get("_dynamic_kv_opt_length", max_cache_length)) + dynamic_kv_opt_rows = max(1, min(dynamic_kv_opt_rows, max_cache_length)) + raw_profile_rows = config.raw.get("_dynamic_kv_profile_rows") + if raw_profile_rows: + dynamic_kv_profile_rows = [] + for row in raw_profile_rows: + clamped = max(1, min(int(row), max_cache_length)) + if clamped not in dynamic_kv_profile_rows: + dynamic_kv_profile_rows.append(clamped) + dynamic_kv_profile_rows.sort() + else: + dynamic_kv_profile_rows = [] + + logger = trt.Logger(trt.Logger.VERBOSE if verbose else trt.Logger.WARNING) + builder = trt.Builder(logger) + network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) + trt_config = builder.create_builder_config() + + # Precision configuration + if precision == "fp16": + work_np_dtype = np.float16 + work_trt_dtype = trt.float16 + elif precision == "bf16": + work_np_dtype = np.float16 # stored as float16, TRT uses bfloat16 + work_trt_dtype = trt.bfloat16 + else: + work_np_dtype = np.float32 + work_trt_dtype = trt.float32 + + if dynamic_kv_cache and position_type == "alibi": + raise ValueError("dynamic_kv_cache is not supported for ALiBi decoder builds") + + # --------------------------------------------------------------- + # Inputs + # --------------------------------------------------------------- + token_id = network.add_input("token_id", trt.int32, (1,)) + position_id = network.add_input("position_id", trt.int32, (1,)) + attention_mask = network.add_input( + "attention_mask", trt.float32, + (1, -1) if dynamic_kv_cache else (1, attention_window)) + + # Optional VL inputs: when embed_input=True, the decoder can accept + # a pre-computed embedding vector instead of a token ID. + input_embed_tensor = None + use_input_embed_tensor = None + if embed_input: + input_embed_tensor = network.add_input( + "input_embed", trt.float32, (1, hidden)) + use_input_embed_tensor = network.add_input( + "use_input_embed", trt.float32, (1,)) + + cache_k_inputs = [] + cache_v_inputs = [] + for i in range(num_layers): + ck = network.add_input( + graph_ops.layer_tensor_name("cache_k", i), + work_trt_dtype, + (-1, kv_attention_size) if dynamic_kv_cache else ( + max_cache_length, kv_attention_size)) + cv = network.add_input( + graph_ops.layer_tensor_name("cache_v", i), + work_trt_dtype, + (-1, kv_attention_size) if dynamic_kv_cache else ( + max_cache_length, kv_attention_size)) + cache_k_inputs.append(ck) + cache_v_inputs.append(cv) + + if dynamic_kv_cache: + if dynamic_kv_profile_rows: + for profile_rows in dynamic_kv_profile_rows: + profile = builder.create_optimization_profile() + # Keep all profiles valid for short prompts / early decode steps. + # The profile-specific value is the opt/max row budget, not a + # lower bound on the live cache length. + min_rows = 1 + profile.set_shape("attention_mask", + (1, min_rows + 1), + (1, profile_rows + 1), + (1, profile_rows + 1)) + for i in range(num_layers): + min_cache_shape = (min_rows, kv_attention_size) + cache_shape = (profile_rows, kv_attention_size) + profile.set_shape(graph_ops.layer_tensor_name("cache_k", i), + min_cache_shape, cache_shape, cache_shape) + profile.set_shape(graph_ops.layer_tensor_name("cache_v", i), + min_cache_shape, cache_shape, cache_shape) + trt_config.add_optimization_profile(profile) + else: + profile = builder.create_optimization_profile() + profile.set_shape("attention_mask", (1, 2), (1, dynamic_kv_opt_rows + 1), + (1, attention_window)) + for i in range(num_layers): + profile.set_shape(graph_ops.layer_tensor_name("cache_k", i), + (1, kv_attention_size), + (dynamic_kv_opt_rows, kv_attention_size), + (max_cache_length, kv_attention_size)) + profile.set_shape(graph_ops.layer_tensor_name("cache_v", i), + (1, kv_attention_size), + (dynamic_kv_opt_rows, kv_attention_size), + (max_cache_length, kv_attention_size)) + trt_config.add_optimization_profile(profile) + + # Cast attention mask to work dtype for elementwise compatibility + if work_trt_dtype != trt.float32: + mask_cast = network.add_cast(attention_mask, work_trt_dtype) + attention_mask = mask_cast.get_output(0) + + def _cast_work_dtype(tensor: trt.ITensor) -> trt.ITensor: + if tensor.dtype == work_trt_dtype: + return tensor + return network.add_cast(tensor, work_trt_dtype).get_output(0) + + # --------------------------------------------------------------- + # Shared constants + # --------------------------------------------------------------- + embedding_table = graph_ops.add_constant( + network, (vocab, hidden), weights["embedding"], dtype=work_np_dtype) + + # RoPE tables (only needed when position_type == "rope") + position_embed_table = None + alibi_slopes_tensor = None + alibi_indices_tensor = None + + # Native RoPE tensors for IRotaryEmbeddingLayer (TRT 10+). + # Shape: [attention_window, rotary_ndims // 2]. + cos_half_tensor = None + sin_half_tensor = None + rotary_embedding_dim = int(head_dim * partial_rotary_factor) + + if position_type == "rope": + graph_ops.validate_native_rope_dim(rotary_embedding_dim) + cos_half_np = graph_ops.make_rope_table_half_dim( + attention_window, head_dim, config.rope_theta, True, + partial_rotary_factor, interleaved=interleaved_rope, + rope_scaling=config.raw.get("rope_scaling")) + sin_half_np = graph_ops.make_rope_table_half_dim( + attention_window, head_dim, config.rope_theta, False, + partial_rotary_factor, interleaved=interleaved_rope, + rope_scaling=config.raw.get("rope_scaling")) + cos_half_tensor = graph_ops.add_constant( + network, cos_half_np.shape, cos_half_np, dtype=work_np_dtype) + cos_half_tensor = _cast_work_dtype(cos_half_tensor) + sin_half_tensor = graph_ops.add_constant( + network, sin_half_np.shape, sin_half_np, dtype=work_np_dtype) + sin_half_tensor = _cast_work_dtype(sin_half_tensor) + elif position_type == "learned": + pos_embed_np = weights["position_embedding"] + position_embed_table = graph_ops.add_constant( + network, pos_embed_np.shape, pos_embed_np, dtype=work_np_dtype) + elif position_type == "alibi": + alibi_slopes_np = graph_ops.compute_alibi_slopes(num_heads) * float(alibi_bias_scale) + alibi_slopes_tensor = graph_ops.add_constant( + network, (num_heads, 1, 1), + alibi_slopes_np.reshape(num_heads, 1, 1), dtype=np.float32) + # Cache position indices [0, 1, ..., max_cache_length-1]. + # The current token's position (position_id) is appended at runtime. + alibi_indices_tensor = graph_ops.add_constant( + network, (max_cache_length,), + np.arange(max_cache_length, dtype=np.float32), + dtype=np.float32) + + eps_tensor = graph_ops.add_constant( + network, (1, 1), np.array([config.rms_norm_eps], dtype=work_np_dtype), + dtype=work_np_dtype) + attn_scale = (1.0 / np.sqrt(max(head_dim, 1))) if scale_attn_weights else 1.0 + # --------------------------------------------------------------- + # Embedding lookup (with optional embed_input override for VL) + # --------------------------------------------------------------- + gather = network.add_gather(embedding_table, token_id, 0) + token_embed = gather.get_output(0) + + if embed_input and input_embed_tensor is not None and use_input_embed_tensor is not None: + # Conditional embedding: (1 - flag) * token_embed + flag * input_embed + # use_input_embed is [1] scalar (FP32), broadcast to [1, hidden] + flag_broadcast = network.add_shuffle(use_input_embed_tensor) + flag_broadcast.reshape_dims = (1, 1) + # Cast flag to work dtype for elementwise compatibility + flag_for_math = flag_broadcast.get_output(0) + if work_trt_dtype != trt.float32: + flag_for_math = network.add_cast(flag_for_math, work_trt_dtype).get_output(0) + one_const = const_in_work_dtype( + network, (1, 1), np.array([1.0], dtype=work_np_dtype), + work_np_dtype, work_trt_dtype) + token_embed = _cast_work_dtype(token_embed) + inv_flag = network.add_elementwise( + one_const, flag_for_math, + trt.ElementWiseOperation.SUB) + # (1 - flag) * token_embed + tok_part = network.add_elementwise( + inv_flag.get_output(0), token_embed, + trt.ElementWiseOperation.PROD) + # flag * input_embed + embed_part = network.add_elementwise( + flag_for_math, _cast_work_dtype(input_embed_tensor), + trt.ElementWiseOperation.PROD) + # sum + hidden_state_sum = network.add_elementwise( + tok_part.get_output(0), embed_part.get_output(0), + trt.ElementWiseOperation.SUM) + hidden_state = hidden_state_sum.get_output(0) + else: + hidden_state = token_embed + + # Add learned position embedding if applicable + if position_type == "learned" and position_embed_table is not None: + pos_gather = network.add_gather(position_embed_table, position_id, 0) + pos_add = network.add_elementwise( + hidden_state, pos_gather.get_output(0), + trt.ElementWiseOperation.SUM) + hidden_state = pos_add.get_output(0) + + # In BF16 mode many embedding/position constants are still materialized from + # float16 storage. Normalize the decoder's main hidden stream back to the + # requested runtime dtype before entering the layer stack. + if hidden_state.dtype != work_trt_dtype: + hidden_state = network.add_cast(hidden_state, work_trt_dtype).get_output(0) + + # Optional embedding LayerNorm (e.g. BLOOM) — use native INormalizationLayer + embed_norm = weights.get("embedding_norm") + if embed_norm is not None: + embed_norm_beta = weights.get("embedding_norm_beta") + if embed_norm_beta is None: + embed_norm_beta = np.zeros(hidden, dtype=work_np_dtype) + hidden_state = graph_ops.add_layer_norm_native( + network, hidden_state, hidden, embed_norm, embed_norm_beta, + config.rms_norm_eps, dtype=work_np_dtype) + + if debug_layer_outputs: + _mark_debug_output(network, hidden_state, "debug_embed") + + # FFI attention kernel: set by the perf agent on their branch. + # Default: None (use native TRT attention). + ffi_attention_kernel = None + + # --------------------------------------------------------------- + # Decoder layers + # --------------------------------------------------------------- + present_k_outputs = [] + present_v_outputs = [] + + # SmolLM3 interleaves NoPE layers; layers flagged False here skip RoPE. + rope_schedule = resolve_rope_layer_schedule(config) + + for layer_idx in range(num_layers): + prefix = f"layer.{layer_idx}" + layer_is_fp32 = layer_idx in fp32_layers + layer_np_dtype = np.float32 if layer_is_fp32 else work_np_dtype + layer_trt_dtype = trt.float32 if layer_is_fp32 else work_trt_dtype + + def _cast_layer_dtype(tensor: trt.ITensor | None) -> trt.ITensor | None: + if tensor is None or tensor.dtype == layer_trt_dtype: + return tensor + return network.add_cast(tensor, layer_trt_dtype).get_output(0) + + result = _add_decoder_layer( + network=network, + hidden=_cast_layer_dtype(hidden_state), + cache_k=_cast_layer_dtype(cache_k_inputs[layer_idx]), + cache_v=_cast_layer_dtype(cache_v_inputs[layer_idx]), + attention_mask=_cast_layer_dtype(attention_mask), + position_id=position_id, + attention_scale=attn_scale, + eps_tensor=_cast_layer_dtype(eps_tensor), + eps=config.rms_norm_eps, + weights=weights, + prefix=prefix, + hidden_size=hidden, + attention_size=attention_size, + kv_attention_size=kv_attention_size, + mlp_size=mlp_size, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + max_cache_length=max_cache_length, + norm_type=norm_type, + mlp_type=mlp_type, + position_type=position_type, + apply_rope=rope_schedule[layer_idx], + activation=activation, + parallel_residual=parallel_residual, + alibi_slopes_tensor=alibi_slopes_tensor, + alibi_indices_tensor=alibi_indices_tensor, + dtype=layer_np_dtype, + quant_ctx=quant_ctx, + cos_half_tensor=_cast_layer_dtype(cos_half_tensor), + sin_half_tensor=_cast_layer_dtype(sin_half_tensor), + rotary_embedding_dim=rotary_embedding_dim, + interleaved_rope=interleaved_rope, + ffi_attention_kernel=ffi_attention_kernel, + dynamic_kv_cache=dynamic_kv_cache, + ) + + hidden_state = _cast_work_dtype(result["hidden"]) + present_k_outputs.append(_cast_work_dtype(result["present_k"])) + present_v_outputs.append(_cast_work_dtype(result["present_v"])) + + if debug_layer_outputs: + _mark_debug_output(network, result["post_attn"], f"debug_post_attn_{layer_idx}") + _mark_debug_output(network, hidden_state, f"debug_hidden_{layer_idx}") + + # --------------------------------------------------------------- + # Final norm + # --------------------------------------------------------------- + final_norm = weights.get("final_norm") + if final_norm is not None and len(final_norm) > 0: + hidden_state = _apply_norm( + network, hidden_state, hidden, final_norm, + weights.get("final_norm_beta"), eps_tensor, norm_type, + dtype=work_np_dtype, eps=config.rms_norm_eps) + + # Optional: mark hidden state as extra output for speech pipelines + if hidden_state_output: + hs_out = network.add_identity(hidden_state).get_output(0) + hs_out.name = "hidden_state" + network.mark_output(hs_out) + + # --------------------------------------------------------------- + # LM head (logits) + # --------------------------------------------------------------- + # Output vocab may differ from input vocab (e.g. Bark semantic: 129600 in, 10048 out). + # Derive from w_out shape if available. + out_vocab = weights["w_out"].shape[1] if isinstance(weights["w_out"], np.ndarray) else vocab + logits = graph_ops.add_matmul_rhs_constant( + network, hidden_state, hidden, out_vocab, weights["w_out"], + dtype=work_np_dtype) + # LM head bias (if present, e.g. CodeGen) or zero bias for C++ parity + lm_bias = weights.get("lm_head_bias") + if lm_bias is not None: + logits = graph_ops.add_bias_sum(network, logits, out_vocab, lm_bias, + dtype=work_np_dtype) + else: + b_out = np.zeros(out_vocab, dtype=work_np_dtype) + logits = graph_ops.add_bias_sum(network, logits, out_vocab, b_out, + dtype=work_np_dtype) + + # Logits output: always FP32 for accurate argmax/sampling + if work_trt_dtype != trt.float32: + logits_cast = network.add_cast(logits, trt.float32) + logits = logits_cast.get_output(0) + logits.name = "logits" + network.mark_output(logits) + + # --------------------------------------------------------------- + # Present K/V outputs + # --------------------------------------------------------------- + for i in range(num_layers): + pk = present_k_outputs[i] + pv = present_v_outputs[i] + pk.name = graph_ops.layer_tensor_name("present_k", i) + pv.name = graph_ops.layer_tensor_name("present_v", i) + network.mark_output(pk) + network.mark_output(pv) + + # --------------------------------------------------------------- + # Build engine + # --------------------------------------------------------------- + if verbose: + print(f"[trtmc build] Building TRT engine ({num_layers} layers, " + f"hidden={hidden}, attn={attention_size}, kv={kv_attention_size}, " + f"mlp={mlp_size}, " + f"cache={max_cache_length}, precision={precision}) ...", + file=sys.stderr) + + plan = builder.build_serialized_network(network, trt_config) + if plan is None: + raise RuntimeError("TensorRT engine build failed") + + return bytes(plan) + + +def _apply_norm( + network: trt.INetworkDefinition, + inp: trt.ITensor, + hidden_size: int, + gamma: np.ndarray, + beta: np.ndarray | None, + eps_tensor: trt.ITensor, + norm_type: str, + dtype: np.dtype = np.float32, + eps: float | None = None, +) -> trt.ITensor: + """Dispatch to RMSNorm or LayerNorm based on norm_type.""" + return graph_blocks.apply_norm( + network, inp, hidden_size, gamma, beta, eps_tensor, norm_type, + dtype=dtype, eps=eps) + + +def _add_decoder_layer( + *, + network: trt.INetworkDefinition, + hidden: trt.ITensor, + cache_k: trt.ITensor, + cache_v: trt.ITensor, + attention_mask: trt.ITensor, + position_id: trt.ITensor, + attention_scale: float | None, + eps_tensor: trt.ITensor, + weights: WeightDict, + prefix: str, + hidden_size: int, + attention_size: int, + kv_attention_size: int, + mlp_size: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + max_cache_length: int, + norm_type: str = "rmsnorm", + mlp_type: str = "swiglu", + position_type: str = "rope", + apply_rope: bool = True, + activation: str = "silu", + parallel_residual: bool = False, + alibi_slopes_tensor: trt.ITensor | None = None, + alibi_indices_tensor: trt.ITensor | None = None, + dtype: np.dtype = np.float32, + quant_ctx: QuantContext | None = None, + cos_half_tensor: trt.ITensor | None = None, + sin_half_tensor: trt.ITensor | None = None, + rotary_embedding_dim: int = 0, + interleaved_rope: bool = False, + ffi_attention_kernel: str | None = None, + dynamic_kv_cache: bool = False, + eps: float | None = None, +) -> dict[str, trt.ITensor]: + """Add one standard decoder layer block. Returns hidden, present_k, present_v.""" + + # Attention block (pre-norm -> QKV -> RoPE -> cache -> attn -> out proj) + attn = graph_blocks.add_attention_block( + network, hidden, cache_k, cache_v, attention_mask, position_id, + weights=weights, prefix=prefix, + hidden_size=hidden_size, attention_size=attention_size, + kv_attention_size=kv_attention_size, + num_heads=num_heads, num_kv_heads=num_kv_heads, head_dim=head_dim, + max_cache_length=max_cache_length, + attention_scale=attention_scale, + eps_tensor=eps_tensor, eps=eps, + norm_type=norm_type, position_type=position_type, + apply_rope=apply_rope, + alibi_slopes_tensor=alibi_slopes_tensor, + alibi_indices_tensor=alibi_indices_tensor, + dtype=dtype, + quant_ctx=quant_ctx, + layer_prefix=prefix, + cos_half_tensor=cos_half_tensor, + sin_half_tensor=sin_half_tensor, + rotary_embedding_dim=rotary_embedding_dim, + interleaved_rope=interleaved_rope, + ffi_attention_kernel=ffi_attention_kernel, + dynamic_kv_cache=dynamic_kv_cache, + ) + attn_out = attn["attn_out"] + present_k = attn["present_k"] + present_v = attn["present_v"] + + # --- Parallel vs sequential residual --- + if parallel_residual: + post_attn_norm_w = weights.get(f"{prefix}.post_attn_norm") + if post_attn_norm_w is not None: + norm2 = _apply_norm( + network, hidden, hidden_size, + post_attn_norm_w, + weights.get(f"{prefix}.post_attn_norm_beta"), + eps_tensor, norm_type, dtype=dtype, eps=eps) + else: + norm2 = attn["normed"] + else: + residual1 = network.add_elementwise( + hidden, attn_out, trt.ElementWiseOperation.SUM) + norm2 = _apply_norm( + network, residual1.get_output(0), hidden_size, + weights[f"{prefix}.post_attn_norm"], + weights.get(f"{prefix}.post_attn_norm_beta"), + eps_tensor, norm_type, dtype=dtype, eps=eps) + + # MLP + if mlp_type == "gelu_fc": + mlp_out = graph_blocks.add_gelu_fc_mlp( + network, norm2, weights=weights, prefix=prefix, + hidden_size=hidden_size, mlp_size=mlp_size, + activation=activation, dtype=dtype, + quant_ctx=quant_ctx, layer_prefix=prefix) + else: + mlp_out = graph_blocks.add_swiglu_mlp( + network, norm2, weights=weights, prefix=prefix, + hidden_size=hidden_size, mlp_size=mlp_size, dtype=dtype, + quant_ctx=quant_ctx, layer_prefix=prefix) + + # Final residual connection + if parallel_residual: + sum_attn = network.add_elementwise( + hidden, attn_out, trt.ElementWiseOperation.SUM) + residual2 = network.add_elementwise( + sum_attn.get_output(0), mlp_out, trt.ElementWiseOperation.SUM) + post_attn_tensor = sum_attn.get_output(0) + else: + residual2 = network.add_elementwise( + residual1.get_output(0), mlp_out, trt.ElementWiseOperation.SUM) + post_attn_tensor = residual1.get_output(0) + + return { + "hidden": residual2.get_output(0), + "post_attn": post_attn_tensor, + "present_k": present_k, + "present_v": present_v, + } diff --git a/python/tensorrt_model_connect/families/smollm3/utils.py b/python/tensorrt_model_connect/families/smollm3/utils.py new file mode 100644 index 0000000000..55986de334 --- /dev/null +++ b/python/tensorrt_model_connect/families/smollm3/utils.py @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Model-agnostic helpers for TensorRT engine builders.""" + +from __future__ import annotations + + +import numpy as np +from tensorrt_model_connect import trt_compat + +from . import graph_ops + + +trt = trt_compat.get_trt() + + +def const_in_work_dtype( + network: trt.INetworkDefinition, + shape: tuple, + values: np.ndarray, + work_np_dtype: np.dtype, + work_trt_dtype: trt.DataType, +) -> trt.ITensor: + """Create a constant in storage dtype and cast it to runtime dtype.""" + const = graph_ops.add_constant(network, shape, values, dtype=work_np_dtype) + if const.dtype != work_trt_dtype: + const = network.add_cast(const, work_trt_dtype).get_output(0) + return const diff --git a/src/runtime/models/smollm3/MODEL.toml b/src/runtime/models/smollm3/MODEL.toml new file mode 100644 index 0000000000..961774cfad --- /dev/null +++ b/src/runtime/models/smollm3/MODEL.toml @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +id = "smollm3" +runtime_library = "libtrtmc_model_smollm3.so" +runtime_plugins = ["plugin.cpp|register_smollm3_plugin"] +runtime_strategies = ["smollm3_decoder_kv_cache"] +runtime_tests = [ + "test_smollm3_plugin_helpers|test_smollm3_plugin_helpers.cpp|_|plugin_helpers.cpp|_", + "test_smollm3_chat_template|test_smollm3_chat_template.cpp|trtmc_model_smollm3|_|_", + "test_smollm3_pipeline|test_smollm3_pipeline.cpp|trtmc_model_smollm3,trtmc_backend_trt|_|REQUIRES_TRT,REQUIRES_GPU", + "test_smollm3_native_kv_cache|test_smollm3_native_kv_cache.cpp|trtmc_model_smollm3|_|REQUIRES_GPU", +] + +[validation_profiles] +decoder_debug = ["smollm3_decoder_kv_cache"] diff --git a/src/runtime/models/smollm3/argmax_kernel.cu b/src/runtime/models/smollm3/argmax_kernel.cu new file mode 100644 index 0000000000..1d51179444 --- /dev/null +++ b/src/runtime/models/smollm3/argmax_kernel.cu @@ -0,0 +1,91 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +// GPU-side argmax for greedy token selection. +// +// Eliminates the D2H transfer of the full logit vector (~600KB for 151K vocab) +// per decode step. Instead, runs a parallel reduction on GPU and copies back +// only the single int32 token ID (4 bytes). +// +// Architecture: single-block parallel reduction using shared memory. +// For vocab_size up to ~256K, one block of 256 threads is sufficient +// (each thread handles ceil(vocab_size/256) elements). + +#include "runtime/models/smollm3/argmax_kernel.h" + +#include +#include + +namespace trtmc { + +static constexpr int kBlockSize = 256; + +__global__ void argmax_reduce_kernel( + const float* __restrict__ logits, + int32_t vocab_size, + int32_t* __restrict__ out_token_id, + float* __restrict__ out_logit) +{ + __shared__ float s_vals[kBlockSize]; + __shared__ int32_t s_idxs[kBlockSize]; + + const int tid = threadIdx.x; + + // Each thread finds the max over its strided range + float best_val = -FLT_MAX; + int32_t best_idx = 0; + + for (int i = tid; i < vocab_size; i += kBlockSize) + { + float v = logits[i]; + if (v > best_val) + { + best_val = v; + best_idx = i; + } + } + + s_vals[tid] = best_val; + s_idxs[tid] = best_idx; + __syncthreads(); + + // Tree reduction + for (int stride = kBlockSize / 2; stride > 0; stride >>= 1) + { + if (tid < stride) + { + if (s_vals[tid + stride] > s_vals[tid]) + { + s_vals[tid] = s_vals[tid + stride]; + s_idxs[tid] = s_idxs[tid + stride]; + } + } + __syncthreads(); + } + + // Thread 0 writes result + if (tid == 0) + { + *out_token_id = s_idxs[0]; + if (out_logit) + { + *out_logit = s_vals[0]; + } + } +} + +void smollm3_gpu_argmax( + const float* d_logits, + int32_t vocab_size, + int32_t* d_token_id, + float* d_logit_val, + cudaStream_t stream) +{ + if (vocab_size <= 0) return; + argmax_reduce_kernel<<<1, kBlockSize, 0, stream>>>( + d_logits, vocab_size, d_token_id, d_logit_val); +} + +} // namespace trtmc diff --git a/src/runtime/models/smollm3/argmax_kernel.h b/src/runtime/models/smollm3/argmax_kernel.h new file mode 100644 index 0000000000..267f136b7d --- /dev/null +++ b/src/runtime/models/smollm3/argmax_kernel.h @@ -0,0 +1,20 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include + +namespace trtmc { + +// GPU-side argmax over a float logit vector. +// Writes the index of the maximum element to d_token_id (device memory). +// Optionally writes the max logit value to d_logit_val (pass nullptr to skip). +// Runs asynchronously on the given stream. +void smollm3_gpu_argmax(const float* d_logits, int32_t vocab_size, int32_t* d_token_id, + float* d_logit_val, cudaStream_t stream); + +} // namespace trtmc diff --git a/src/runtime/models/smollm3/chat_templates.cpp b/src/runtime/models/smollm3/chat_templates.cpp new file mode 100644 index 0000000000..427c602759 --- /dev/null +++ b/src/runtime/models/smollm3/chat_templates.cpp @@ -0,0 +1,169 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/models/smollm3/chat_templates.h" + +#include +#include + +namespace trtmc { +namespace { + +std::string apply_chatml(const std::string& prompt, bool enable_thinking) { + std::string r = "<|im_start|>user\n" + prompt + "<|im_end|>\n<|im_start|>assistant\n"; + if (!enable_thinking) + r += "\n\n\n\n"; + return r; +} + +std::string apply_mistral(const std::string& prompt, bool /*enable_thinking*/) { + return "[INST] " + prompt + " [/INST]"; +} + +std::string apply_phi(const std::string& prompt, bool /*enable_thinking*/) { + return "<|user|>\n" + prompt + "<|end|>\n<|assistant|>\n"; +} + +std::string apply_gemma(const std::string& prompt, bool /*enable_thinking*/) { + return "user\n" + prompt + "\nmodel\n"; +} + +std::string apply_llama3(const std::string& prompt, bool enable_thinking) { + std::string r = "<|begin_of_text|>"; + if (!enable_thinking) + r += "<|start_header_id|>system<|end_header_id|>\n\ndetailed thinking off<|eot_id|>"; + r += "<|start_header_id|>user<|end_header_id|>\n\n" + prompt + + "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"; + return r; +} + +std::string apply_nemotron(const std::string& prompt, bool /*enable_thinking*/) { + return "System\n\nUser\n" + prompt + "\nAssistant\n"; +} + +std::string apply_nemotron_h(const std::string& prompt, bool enable_thinking) { + std::string r = + "System\n\nUser\n" + prompt + "\nAssistant\n"; + r += enable_thinking ? "\n" : ""; + return r; +} + +// SmolLM3 always emits a system block, even when the caller supplies no system +// message. The two Custom Instruction bodies below are the upstream defaults for +// the two reasoning modes; reproducing them verbatim is what keeps the served +// prompt identical to the Hugging Face chat template. +constexpr char kSmollm3ThinkInstructions[] = + "You are a helpful AI assistant named SmolLM, trained by Hugging Face. Your role as an " + "assistant involves thoroughly exploring questions through a systematic thinking process " + "before providing the final precise and accurate solutions. This requires engaging in a " + "comprehensive cycle of analysis, summarizing, exploration, reassessment, reflection, " + "backtracking, and iteration to develop well-considered thinking process. Please structure " + "your response into two main sections: Thought and Solution using the specified format: " + " Thought section Solution section. In the Thought section, detail your " + "reasoning process in steps. Each step should include detailed considerations such as " + "analysing questions, summarizing relevant findings, brainstorming new ideas, verifying the " + "accuracy of the current steps, refining any errors, and revisiting previous steps. In the " + "Solution section, based on various attempts, explorations, and reflections from the Thought " + "section, systematically present the final solution that you deem correct. The Solution " + "section should be logical, accurate, and concise and detail necessary steps needed to reach " + "the conclusion."; +constexpr char kSmollm3NoThinkInstructions[] = + "You are a helpful AI assistant named SmolLM, trained by Hugging Face."; + +std::string smollm3_today() { + // Upstream renders strftime_now("%d %B %Y"), e.g. "04 September 2026". + std::time_t now = std::time(nullptr); + std::tm tm_utc{}; +#if defined(_WIN32) + gmtime_s(&tm_utc, &now); +#else + gmtime_r(&now, &tm_utc); +#endif + char buf[64]; + if (std::strftime(buf, sizeof(buf), "%d %B %Y", &tm_utc) == 0) + return {}; + return std::string(buf); +} + +std::string apply_smollm3(const std::string& prompt, bool enable_thinking, + const std::string& today) { + const std::string mode = enable_thinking ? "/think" : "/no_think"; + std::string r = "<|im_start|>system\n"; + r += "## Metadata\n\n"; + r += "Knowledge Cutoff Date: June 2025\n"; + r += "Today Date: " + (today.empty() ? smollm3_today() : today) + "\n"; + r += "Reasoning Mode: " + mode + "\n\n"; + r += "## Custom Instructions\n\n"; + r += enable_thinking ? kSmollm3ThinkInstructions : kSmollm3NoThinkInstructions; + r += "\n\n"; + // Upstream does NOT close the system block with <|im_end|>: the instructions + // run straight into the user turn. Verified against apply_chat_template. + r += "<|im_start|>user\n" + prompt + "<|im_end|>\n"; + r += "<|im_start|>assistant\n"; + if (!enable_thinking) + r += "\n\n\n"; + return r; +} + +} // namespace + +std::string smollm3_detect_chat_template_format(const std::string& jinja_template) { + // Ordered marker table: the first row whose markers are all present wins. + // SmolLM3's template is ChatML-framed, so it is matched on its mandatory + // system block before the generic ChatML row below. A row's second marker + // is either an additional requirement (`require_both`) or an alternative. + struct Rule { + const char* first; + const char* second; + bool require_both; + const char* format; + }; + static constexpr Rule kRules[] = { + {"<|im_start|>", "Reasoning Mode:", true, "smollm3"}, + {"<|im_start|>", nullptr, false, "chatml"}, + {"[INST]", nullptr, false, "mistral"}, + {"<|user|>", "<|assistant|>", false, "phi"}, + {"", nullptr, false, "gemma"}, + {"<|start_header_id|>", nullptr, false, "llama3"}, + {"", nullptr, false, "nemotron"}, + {"", nullptr, false, "nemotron_h"}, + }; + + for (const Rule& rule : kRules) { + const bool has_first = jinja_template.find(rule.first) != std::string::npos; + const bool has_second = + rule.second != nullptr && jinja_template.find(rule.second) != std::string::npos; + const bool matched = + rule.require_both ? (has_first && has_second) : (has_first || has_second); + if (matched) + return rule.format; + } + return {}; +} + +std::string smollm3_apply_chat_template(const std::string& format, const std::string& prompt, + bool enable_thinking, const std::string& today) { + if (format.empty()) + return prompt; + if (format == "smollm3") + return apply_smollm3(prompt, enable_thinking, today); + if (format == "chatml") + return apply_chatml(prompt, enable_thinking); + if (format == "mistral") + return apply_mistral(prompt, enable_thinking); + if (format == "phi") + return apply_phi(prompt, enable_thinking); + if (format == "gemma") + return apply_gemma(prompt, enable_thinking); + if (format == "llama3") + return apply_llama3(prompt, enable_thinking); + if (format == "nemotron") + return apply_nemotron(prompt, enable_thinking); + if (format == "nemotron_h") + return apply_nemotron_h(prompt, enable_thinking); + return prompt; +} + +} // namespace trtmc diff --git a/src/runtime/models/smollm3/chat_templates.h b/src/runtime/models/smollm3/chat_templates.h new file mode 100644 index 0000000000..cd8f136cd8 --- /dev/null +++ b/src/runtime/models/smollm3/chat_templates.h @@ -0,0 +1,16 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +namespace trtmc { + +std::string smollm3_detect_chat_template_format(const std::string& jinja_template); +std::string smollm3_apply_chat_template(const std::string& format, const std::string& prompt, + bool enable_thinking = true, const std::string& today = {}); + +} // namespace trtmc diff --git a/src/runtime/models/smollm3/inference_state.h b/src/runtime/models/smollm3/inference_state.h new file mode 100644 index 0000000000..a6ab14f28f --- /dev/null +++ b/src/runtime/models/smollm3/inference_state.h @@ -0,0 +1,104 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +// Smollm3InferenceState: unified interface for autoregressive inference state. +// +// Both KV-cache attention state and recurrent state implementations expose +// this interface. Pipelines and plugins program against it — never against +// concrete state classes. +// +// The interface captures the lifecycle of per-sequence inference state: +// 1. reset() — prepare for a new sequence +// 2. bind_to() — bind state tensors to TRT engine I/O +// 3. prepare_step() — write state-related inputs (mask, position) into TensorMap +// 4. advance() — update state after each decode step +// 5. position() — current sequence position +// +// Implementations: +// Smollm3KvCache — dense append-only (current default) +// Family-owned recurrent state — recurrent tensor state +// Family-owned hybrid state — Smollm3KvCache + family-owned recurrent state composed +// (future: RingKvCache, PagedKvCache, MlaCache, SlidingWindowCache) + +#include "trtmc/runtime/tensor.h" + +#include +#include + +namespace trtmc { + +class ITrtModule; +using TrtModule = ITrtModule; + +class Smollm3InferenceState { + public: + virtual ~Smollm3InferenceState() = default; + + // --- Lifecycle --- + + // Reset logical state for a new sequence. Implementations may retain device + // storage that remains hidden by logical lengths and attention masks. + virtual void reset() = 0; + + // Bind all state tensors to the given TRT module. + // Called once per sequence after reset(). The module reads/writes + // state tensors via the bound device pointers. + virtual void bind_to(TrtModule& module) = 0; + + // Write state-related inputs (mask, position, block table, etc.) into + // the TensorMap before engine.forward(). The state owns its buffers — + // Tensor.data pointers remain valid until the next prepare_step() call. + // Pipelines call this instead of manually constructing mask/position tensors. + virtual void prepare_step(TensorMap& inputs, int32_t seq_len = 1) = 0; + + // Update state after one decode step. Copies "present" outputs + // into "cache" inputs, advances position. + // n_tokens: number of tokens processed in this step (default 1). + // >1 for batched prefill / multi-token steps. + virtual void advance(int32_t n_tokens = 1) = 0; + + // Provide the total prompt length before prefill starts. + // Cache policies that distinguish prompt and decode tokens can use this + // to protect prompt tokens even if compression triggers during prefill. + virtual void set_prompt_length(int32_t prompt_length) { (void)prompt_length; } + + // Mark the transition from prompt prefill to autoregressive decoding. + // State types that do not distinguish the phases can ignore this. + virtual void mark_prefill_complete() {} + + // --- Queries --- + + // Current sequence position (0 = empty, increments with advance()). + virtual int32_t position() const = 0; + + // Maximum sequence length this state can hold. + // -1 for unbounded (recurrent models with no cache length limit). + virtual int32_t max_length() const = 0; + + // Desired number of KV rows to expose to the decoder on the next step. + // Dynamic-KV runtimes can use this to choose an execution profile/context + // before prepare_step() binds the state tensors. + virtual int32_t preferred_cache_rows() const { return max_length(); } + + // Number of transformer/SSM layers. + virtual int32_t num_layers() const = 0; + + // Whether this state type needs an attention mask. + // Smollm3KvCache -> true. Family-owned recurrent state -> false. + virtual bool needs_attention_mask() const = 0; + + // Total device memory consumed by this state (bytes). + virtual std::size_t device_memory_bytes() const = 0; + + // Human-readable state type for diagnostics. + virtual const char* state_type() const = 0; + + // Whether all allocations succeeded. + virtual bool ok() const = 0; +}; + +} // namespace trtmc diff --git a/src/runtime/models/smollm3/kv_cache.cpp b/src/runtime/models/smollm3/kv_cache.cpp new file mode 100644 index 0000000000..731780e482 --- /dev/null +++ b/src/runtime/models/smollm3/kv_cache.cpp @@ -0,0 +1,568 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/models/smollm3/kv_cache.h" + +#include "trtmc/runtime/trt_module.h" + +#include +#include +#include +#include + +namespace trtmc { + +namespace { + +constexpr int32_t kRuntimeBucketRows = 32; + +int32_t round_up_rows(int32_t value, int32_t bucket, int32_t maximum) { + if (bucket <= 1) + return std::min(std::max(value, 1), maximum); + const int32_t rounded = ((std::max(value, 1) + bucket - 1) / bucket) * bucket; + return std::min(rounded, maximum); +} + +void validate_native_scalar_input(TrtModule& module, const std::string& name) { + if (module.tensor_dtype(name) != DType::kInt32 || + module.tensor_shape(name) != std::vector{1}) { + throw std::runtime_error("SmolLM3 native KV input '" + name + "' must be int32 [1]"); + } +} + +bool valid_native_cache_shape(const std::vector& shape, int32_t max_length, + int32_t kv_dim) { + if (shape.size() != 4) + return false; + if (shape[0] != 1 || shape[2] != static_cast(max_length)) + return false; + if (shape[1] <= 0 || shape[3] <= 0) + return false; + return shape[1] * shape[3] == kv_dim; +} + +void validate_native_cache_pair(TrtModule& module, const std::string& cache_name, + const std::string& present_name, int32_t max_length, int32_t kv_dim, + DType cache_dtype) { + if (!module.has_input(cache_name) || !module.has_output(present_name)) { + throw std::runtime_error("SmolLM3 native KV engine is missing cache/present pair '" + + cache_name + "'/'" + present_name + "'"); + } + const auto cache_shape = module.tensor_shape(cache_name); + const auto present_shape = module.tensor_shape(present_name); + if (!valid_native_cache_shape(cache_shape, max_length, kv_dim) || + present_shape != cache_shape) { + throw std::runtime_error("SmolLM3 native KV cache/present tensors must share static " + "[1,Hkv,max_length,D] shape"); + } + if (module.tensor_dtype(cache_name) != cache_dtype || + module.tensor_dtype(present_name) != cache_dtype) { + throw std::runtime_error("SmolLM3 native KV cache dtype does not match model precision"); + } +} + +bool all_tensors_ok(const std::vector& tensors) { + for (const auto& tensor : tensors) { + if (!tensor.ok()) + return false; + } + return true; +} + +} // namespace + +Smollm3KvCache::Smollm3KvCache(int32_t num_layers, int32_t max_length, int32_t kv_dim, + cudaStream_t stream, DType cache_dtype, Smollm3KvCacheNames names) + : num_layers_(num_layers), max_length_(max_length), kv_dim_(kv_dim), stream_(stream), + cache_dtype_(cache_dtype), cache_element_size_(dtype_size(cache_dtype)), + names_(std::move(names)) { + + // If names were not supplied, generate standard defaults. + if (names_.cache_k.empty()) { + names_.cache_k.reserve(static_cast(num_layers)); + names_.cache_v.reserve(static_cast(num_layers)); + names_.present_k.reserve(static_cast(num_layers)); + names_.present_v.reserve(static_cast(num_layers)); + for (int32_t i = 0; i < num_layers; ++i) { + std::string suffix = "_" + std::to_string(i); + names_.cache_k.push_back("cache_k" + suffix); + names_.cache_v.push_back("cache_v" + suffix); + names_.present_k.push_back("present_k" + suffix); + names_.present_v.push_back("present_v" + suffix); + } + } + const auto expected_names = static_cast(num_layers); + if (names_.cache_k.size() != expected_names || names_.cache_v.size() != expected_names || + names_.present_k.size() != expected_names || names_.present_v.size() != expected_names) { + throw std::invalid_argument("Smollm3KvCache: per-layer tensor name count mismatch"); + } + + cache_k_.reserve(static_cast(num_layers)); + cache_v_.reserve(static_cast(num_layers)); + present_k_.reserve(static_cast(num_layers)); + present_v_.reserve(static_cast(num_layers)); + + for (int32_t i = 0; i < num_layers; ++i) { + cache_k_.emplace_back(std::vector{max_length, kv_dim}, cache_dtype_, stream); + if (!cache_k_.back().ok()) + return; + cache_v_.emplace_back(std::vector{max_length, kv_dim}, cache_dtype_, stream); + if (!cache_v_.back().ok()) + return; + } + + // Pre-allocate mask buffer: [max_length + 1] for dense causal mask. + mask_buf_.resize(static_cast(max_length) + 1); + + reset(); +} + +bool Smollm3KvCache::configure_binding_mode(TrtModule& module) { + const bool has_write_indices = module.has_input(names_.cache_write_indices); + const bool has_kv_lengths = module.has_input(names_.key_value_lengths); + if (has_write_indices != has_kv_lengths) { + throw std::runtime_error( + "SmolLM3 native KV engine must expose both cache_write_indices and " + "key_value_lengths"); + } + + const bool native_mode = has_write_indices && has_kv_lengths; + if (binding_mode_initialized_ && native_mode != native_kv_update_enabled_) { + throw std::runtime_error( + "SmolLM3 prefill and decode engines use incompatible KV cache contracts"); + } + binding_mode_initialized_ = true; + native_kv_update_enabled_ = native_mode; + if (native_mode) + validate_native_kv_contract(module); + return native_mode; +} + +void Smollm3KvCache::validate_native_kv_contract(TrtModule& module) const { + validate_native_scalar_input(module, names_.cache_write_indices); + validate_native_scalar_input(module, names_.key_value_lengths); + + for (int32_t i = 0; i < num_layers_; ++i) { + const auto li = static_cast(i); + validate_native_cache_pair(module, names_.cache_k[li], names_.present_k[li], max_length_, + kv_dim_, cache_dtype_); + validate_native_cache_pair(module, names_.cache_v[li], names_.present_v[li], max_length_, + kv_dim_, cache_dtype_); + } +} + +void Smollm3KvCache::ensure_legacy_present_buffers() { + if (!present_k_.empty()) + return; + for (int32_t i = 0; i < num_layers_; ++i) { + present_k_.emplace_back(std::vector{1, kv_dim_}, cache_dtype_, stream_); + present_v_.emplace_back(std::vector{1, kv_dim_}, cache_dtype_, stream_); + if (!present_k_.back().ok() || !present_v_.back().ok()) { + throw std::runtime_error("SmolLM3 legacy KV present-buffer allocation failed"); + } + } +} + +void Smollm3KvCache::bind_native_cache(TrtModule& module) { + for (int32_t i = 0; i < num_layers_; ++i) { + const auto li = static_cast(i); + module.bind_external(names_.cache_k[li], cache_k_[li].data()); + module.bind_external(names_.cache_v[li], cache_v_[li].data()); + if (module.device_ptr(names_.cache_k[li]) != cache_k_[li].data() || + module.device_ptr(names_.present_k[li]) != cache_k_[li].data() || + module.device_ptr(names_.cache_v[li]) != cache_v_[li].data() || + module.device_ptr(names_.present_v[li]) != cache_v_[li].data()) { + throw std::runtime_error( + "SmolLM3 native KV engine did not preserve cache/present aliasing"); + } + } +} + +void Smollm3KvCache::validate_native_aliases(const std::vector& present_k, + const std::vector& present_v) const { + if (static_cast(present_k.size()) != num_layers_ || + static_cast(present_v.size()) != num_layers_) { + throw std::runtime_error("SmolLM3 native KV per-layer pointer count mismatch"); + } + for (int32_t i = 0; i < num_layers_; ++i) { + const auto li = static_cast(i); + if (present_k[li] != cache_k_[li].data() || present_v[li] != cache_v_[li].data()) { + throw std::runtime_error( + "SmolLM3 native prefill present tensors must alias the KV cache"); + } + } +} + +void Smollm3KvCache::write_native_kv_inputs(TensorMap& inputs, int32_t seq_len) { + if (seq_len > max_length_ - position_) { + throw std::runtime_error("SmolLM3 sequence exceeds the model's fixed KV cache capacity"); + } + cache_write_index_ = position_; + key_value_length_ = position_ + seq_len; + inputs[names_.cache_write_indices] = Tensor{&cache_write_index_, {1}, DType::kInt32}; + inputs[names_.key_value_lengths] = Tensor{&key_value_length_, {1}, DType::kInt32}; +} + +// Masked score constant is model-local. +static constexpr float kMaskedScore = -1.0e4F; + +void Smollm3KvCache::build_attention_mask(std::vector& mask) const { + // DEPRECATED: use prepare_step() instead. + const auto width = static_cast(max_length_) + 1; + mask.assign(width, kMaskedScore); + const int32_t valid = std::max(0, std::min(position_, max_length_)); + for (int32_t i = 0; i < valid; ++i) + mask[static_cast(i)] = 0.0f; + mask.back() = 0.0f; +} + +int32_t Smollm3KvCache::preferred_cache_rows() const { + if (!dynamic_binding_enabled_) + return max_length_; + return round_up_rows(std::max(position_, 1), kRuntimeBucketRows, max_length_); +} + +void Smollm3KvCache::rebind_cache_rows(int32_t cache_rows) { + if (!dynamic_binding_enabled_ || bound_module_ == nullptr || cache_rows == bound_cache_rows_) + return; + const std::vector cache_shape{cache_rows, kv_dim_}; + for (int32_t i = 0; i < num_layers_; ++i) { + const auto li = static_cast(i); + bound_module_->bind_external(names_.cache_k[li], cache_k_[li].data(), cache_shape); + bound_module_->bind_external(names_.cache_v[li], cache_v_[li].data(), cache_shape); + } + bound_cache_rows_ = cache_rows; +} + +// Match the engine-declared rank for attention_mask. Different engine families +// wire the causal mask with different shapes: +// * static decoder (cache-full, e.g. legacy builds): [max_length + 1] +// * dynamic decoder (standard + triattention): [1, mask_width] +// * 3-D decoder mask with query dim: [1, 1, mask_width] +// The tensor content is identical (width = current mask_width); only the +// leading broadcast dimensions change. +std::vector Smollm3KvCache::mask_shape_for_engine(int32_t mask_width, + std::size_t mask_buf_size) const { + const int32_t mask_rank = + bound_module_ != nullptr ? bound_module_->input_rank(names_.attention_mask) : 0; + if (mask_rank == 3) + return {1, 1, mask_width}; + if (mask_rank == 2 || (mask_rank == 0 && dynamic_binding_enabled_)) + return {1, mask_width}; + return {static_cast(mask_buf_size)}; +} + +void Smollm3KvCache::write_position_input(TensorMap& inputs, int32_t seq_len) { + if (!has_position_input_) + return; + pos_buf_vec_.resize(static_cast(seq_len)); + for (int32_t i = 0; i < seq_len; ++i) + pos_buf_vec_[static_cast(i)] = position_ + i; + Tensor pos_t; + pos_t.data = pos_buf_vec_.data(); + pos_t.shape = {static_cast(seq_len)}; + pos_t.dtype = DType::kInt32; + inputs[names_.position_id] = pos_t; +} + +void Smollm3KvCache::write_batched_mask(TensorMap& inputs, int32_t seq_len) { + // Batched prefill mask: (seq_len, max_length + seq_len). Columns + // [0, valid) are visible cache, [valid, max_length) are stale slots, + // [max_length, max_length+seq_len) are the new tokens — causal so + // token i sees tokens 0..i. + const int32_t valid = std::max(0, std::min(position_, max_length_)); + const int32_t kv_len = max_length_ + seq_len; + const std::size_t total = static_cast(seq_len) * static_cast(kv_len); + mask_buf_.assign(total, kMaskedScore); + for (int32_t i = 0; i < seq_len; ++i) { + const std::size_t row = static_cast(i) * static_cast(kv_len); + for (int32_t j = 0; j < valid; ++j) + mask_buf_[row + static_cast(j)] = 0.0f; + for (int32_t j = 0; j <= i; ++j) + mask_buf_[row + static_cast(max_length_) + static_cast(j)] = + 0.0f; + } + Tensor mask_t; + mask_t.data = mask_buf_.data(); + mask_t.shape = {static_cast(seq_len), static_cast(kv_len)}; + mask_t.dtype = DType::kFloat32; + inputs[names_.attention_mask] = mask_t; +} + +void Smollm3KvCache::write_bidirectional_mask(TensorMap& inputs, int32_t seq_len) { + // Diffusion block mask: all valid prefix cache rows are visible, stale cache + // rows are hidden, and every token in the current block can see every other + // token in the current block. + const int32_t valid = std::max(0, std::min(position_, max_length_)); + const int32_t kv_len = max_length_ + seq_len; + const std::size_t total = static_cast(seq_len) * static_cast(kv_len); + mask_buf_.assign(total, kMaskedScore); + for (int32_t i = 0; i < seq_len; ++i) { + const std::size_t row = static_cast(i) * static_cast(kv_len); + for (int32_t j = 0; j < valid; ++j) + mask_buf_[row + static_cast(j)] = 0.0f; + for (int32_t j = 0; j < seq_len; ++j) + mask_buf_[row + static_cast(max_length_) + static_cast(j)] = + 0.0f; + } + Tensor mask_t; + mask_t.data = mask_buf_.data(); + mask_t.shape = {static_cast(seq_len), static_cast(kv_len)}; + mask_t.dtype = DType::kFloat32; + inputs[names_.attention_mask] = mask_t; +} + +void Smollm3KvCache::write_decode_mask(TensorMap& inputs) { + const int32_t valid = std::max(0, std::min(position_, max_length_)); + const int32_t cache_rows = dynamic_binding_enabled_ ? preferred_cache_rows() : max_length_; + const int32_t mask_width = dynamic_binding_enabled_ ? (cache_rows + 1) : (max_length_ + 1); + rebind_cache_rows(cache_rows); + + if (mask_buf_.size() < static_cast(mask_width)) + mask_buf_.assign(static_cast(mask_width), kMaskedScore); + std::fill(mask_buf_.begin(), mask_buf_.begin() + mask_width, kMaskedScore); + for (int32_t i = 0; i < valid; ++i) + mask_buf_[static_cast(i)] = 0.0f; + mask_buf_[static_cast(mask_width - 1)] = 0.0f; + + Tensor mask_t; + mask_t.data = mask_buf_.data(); + mask_t.shape = mask_shape_for_engine(mask_width, mask_buf_.size()); + mask_t.dtype = DType::kFloat32; + inputs[names_.attention_mask] = mask_t; +} + +void Smollm3KvCache::prepare_step(TensorMap& inputs, int32_t seq_len) { + if (seq_len <= 0) + seq_len = 1; + write_position_input(inputs, seq_len); + if (native_kv_update_enabled_) { + write_native_kv_inputs(inputs, seq_len); + return; + } + if (seq_len > 1) + write_batched_mask(inputs, seq_len); + else + write_decode_mask(inputs); +} + +void Smollm3KvCache::prepare_bidirectional_step(TensorMap& inputs, int32_t seq_len) { + if (seq_len <= 0) + seq_len = 1; + if (native_kv_update_enabled_) { + throw std::runtime_error("SmolLM3 native TensorRT KV cache supports causal attention only; " + "bidirectional block decoding is unsupported"); + } + write_position_input(inputs, seq_len); + write_bidirectional_mask(inputs, seq_len); +} + +void Smollm3KvCache::bind_to(TrtModule& module) { + bound_module_ = &module; + has_position_input_ = module.has_input(names_.position_id); + if (configure_binding_mode(module)) { + dynamic_binding_enabled_ = false; + bound_cache_rows_ = max_length_; + bind_native_cache(module); + return; + } + ensure_legacy_present_buffers(); + // Enable dynamic row binding only when cache_k[0] itself is dynamic. + // Static-shape engines with fixed [max_length, kv_dim] cache reject + // setInputShape on cache inputs even when other inputs are dynamic. + dynamic_binding_enabled_ = + !names_.cache_k.empty() && module.input_is_dynamic(names_.cache_k.front()); + bound_cache_rows_ = 0; + const int32_t initial_cache_rows = + dynamic_binding_enabled_ ? preferred_cache_rows() : max_length_; + const std::vector cache_shape{initial_cache_rows, kv_dim_}; + + for (int32_t i = 0; i < num_layers_; ++i) { + auto li = static_cast(i); + if (dynamic_binding_enabled_) { + module.bind_external(names_.cache_k[li], cache_k_[li].data(), cache_shape); + module.bind_external(names_.cache_v[li], cache_v_[li].data(), cache_shape); + bound_cache_rows_ = initial_cache_rows; + } else { + module.bind_external(names_.cache_k[li], cache_k_[li].data()); + module.bind_external(names_.cache_v[li], cache_v_[li].data()); + } + module.bind_external(names_.present_k[li], present_k_[li].data()); + module.bind_external(names_.present_v[li], present_v_[li].data()); + } +} + +void Smollm3KvCache::bind_cache_inputs(TrtModule& module) { + bound_module_ = &module; + has_position_input_ = module.has_input(names_.position_id); + if (configure_binding_mode(module)) { + dynamic_binding_enabled_ = false; + bound_cache_rows_ = max_length_; + bind_native_cache(module); + return; + } + // A dual-profile prefill engine exposes dynamic cache rows even though it + // consumes the complete runtime allocation. Bind the runtime-sized view, + // which can be smaller than the engine profile's opt/max shape after a + // --kv-cache-size override. + dynamic_binding_enabled_ = + !names_.cache_k.empty() && module.input_is_dynamic(names_.cache_k.front()); + bound_cache_rows_ = 0; + const std::vector cache_shape{max_length_, kv_dim_}; + for (int32_t i = 0; i < num_layers_; ++i) { + auto li = static_cast(i); + if (dynamic_binding_enabled_) { + module.bind_external(names_.cache_k[li], cache_k_[li].data(), cache_shape); + module.bind_external(names_.cache_v[li], cache_v_[li].data(), cache_shape); + bound_cache_rows_ = max_length_; + } else { + module.bind_external(names_.cache_k[li], cache_k_[li].data()); + module.bind_external(names_.cache_v[li], cache_v_[li].data()); + } + } +} + +void Smollm3KvCache::write_prefill_kv(const std::vector& prefill_k, + const std::vector& prefill_v, int32_t seq_len) { + if (seq_len <= 0) + return; + if (seq_len > max_length_) + throw std::runtime_error("Smollm3KvCache::write_prefill_kv: seq_len exceeds max_length"); + if (static_cast(prefill_k.size()) != num_layers_ || + static_cast(prefill_v.size()) != num_layers_) { + throw std::runtime_error( + "Smollm3KvCache::write_prefill_kv: per-layer pointer count mismatch"); + } + if (native_kv_update_enabled_) { + validate_native_aliases(prefill_k, prefill_v); + position_ = seq_len; + return; + } + const auto row_bytes = static_cast(kv_dim_) * cache_element_size_; + const auto block_bytes = static_cast(seq_len) * row_bytes; + for (int32_t i = 0; i < num_layers_; ++i) { + auto li = static_cast(i); + cudaMemcpyAsync(cache_k_[li].data(), prefill_k[li], block_bytes, cudaMemcpyDeviceToDevice, + stream_); + cudaMemcpyAsync(cache_v_[li].data(), prefill_v[li], block_bytes, cudaMemcpyDeviceToDevice, + stream_); + } + position_ = seq_len; +} + +void Smollm3KvCache::append_prefill_kv(const std::vector& prefill_k, + const std::vector& prefill_v, int32_t seq_len) { + if (seq_len <= 0) + return; + if (position_ + seq_len > max_length_) + throw std::runtime_error("Smollm3KvCache::append_prefill_kv: append exceeds max_length"); + if (static_cast(prefill_k.size()) != num_layers_ || + static_cast(prefill_v.size()) != num_layers_) { + throw std::runtime_error( + "Smollm3KvCache::append_prefill_kv: per-layer pointer count mismatch"); + } + if (native_kv_update_enabled_) { + validate_native_aliases(prefill_k, prefill_v); + position_ += seq_len; + return; + } + const auto row_bytes = static_cast(kv_dim_) * cache_element_size_; + const auto block_bytes = static_cast(seq_len) * row_bytes; + const auto offset = static_cast(position_) * row_bytes; + for (int32_t i = 0; i < num_layers_; ++i) { + auto li = static_cast(i); + cudaMemcpyAsync(static_cast(cache_k_[li].data()) + offset, prefill_k[li], + block_bytes, cudaMemcpyDeviceToDevice, stream_); + cudaMemcpyAsync(static_cast(cache_v_[li].data()) + offset, prefill_v[li], + block_bytes, cudaMemcpyDeviceToDevice, stream_); + } + position_ += seq_len; +} + +void Smollm3KvCache::set_position(int32_t position) { + position_ = std::max(0, std::min(position, max_length_)); +} + +void Smollm3KvCache::advance(int32_t n_tokens) { + // For now, only single-token advance is supported. + // n_tokens > 1 reserved for future batched prefill (TASK-10). + assert(n_tokens == 1 && "Smollm3KvCache::advance: only n_tokens==1 supported"); + (void)n_tokens; + + if (native_kv_update_enabled_) { + if (position_ >= max_length_) { + throw std::runtime_error( + "SmolLM3 sequence exceeds the model's fixed KV cache capacity"); + } + ++position_; + return; + } + + // Copy present K/V (single row) into cache at current position. + // present_k_[layer] is [1, kv_dim] → copy to cache_k_[layer][position_, :] + auto row_bytes = static_cast(kv_dim_) * cache_element_size_; + + if (position_ < max_length_) { + // Normal append: write to position_ slot + auto offset = static_cast(position_) * row_bytes; + for (int32_t i = 0; i < num_layers_; ++i) { + auto li = static_cast(i); + cudaMemcpyAsync(static_cast(cache_k_[li].data()) + offset, + present_k_[li].data(), row_bytes, cudaMemcpyDeviceToDevice, stream_); + cudaMemcpyAsync(static_cast(cache_v_[li].data()) + offset, + present_v_[li].data(), row_bytes, cudaMemcpyDeviceToDevice, stream_); + } + ++position_; + } else { + // Cache full: shift [1..max) → [0..max-1), then write at tail + auto shift_bytes = static_cast(max_length_ - 1) * row_bytes; + auto tail_offset = shift_bytes; + for (int32_t i = 0; i < num_layers_; ++i) { + auto li = static_cast(i); + auto* ck = static_cast(cache_k_[li].data()); + auto* cv = static_cast(cache_v_[li].data()); + cudaMemcpyAsync(ck, ck + row_bytes, shift_bytes, cudaMemcpyDeviceToDevice, stream_); + cudaMemcpyAsync(cv, cv + row_bytes, shift_bytes, cudaMemcpyDeviceToDevice, stream_); + cudaMemcpyAsync(ck + tail_offset, present_k_[li].data(), row_bytes, + cudaMemcpyDeviceToDevice, stream_); + cudaMemcpyAsync(cv + tail_offset, present_v_[li].data(), row_bytes, + cudaMemcpyDeviceToDevice, stream_); + } + // position_ stays at max_length_ (cache is full, all slots visible) + } +} + +void Smollm3KvCache::reset() { + // Reset only the logical sequence length. Attention masks hide every + // stale cache row, and each present row is overwritten before use. + position_ = 0; + cache_write_index_ = 0; + key_value_length_ = 0; +} + +std::size_t Smollm3KvCache::device_memory_bytes() const { + std::size_t total = 0; + for (const auto& t : cache_k_) + total += t.nbytes(); + for (const auto& t : cache_v_) + total += t.nbytes(); + for (const auto& t : present_k_) + total += t.nbytes(); + for (const auto& t : present_v_) + total += t.nbytes(); + return total; +} + +bool Smollm3KvCache::ok() const { + const auto expected_layers = static_cast(num_layers_); + if (cache_k_.size() != expected_layers) + return false; + if (cache_v_.size() != expected_layers) + return false; + return all_tensors_ok(cache_k_) && all_tensors_ok(cache_v_) && all_tensors_ok(present_k_) && + all_tensors_ok(present_v_); +} + +} // namespace trtmc diff --git a/src/runtime/models/smollm3/kv_cache.h b/src/runtime/models/smollm3/kv_cache.h new file mode 100644 index 0000000000..7dcd7da0db --- /dev/null +++ b/src/runtime/models/smollm3/kv_cache.h @@ -0,0 +1,137 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +// Smollm3KvCache: autoregressive KV cache state manager. +// HF equivalent: DynamicCache / past_key_values. +// +// Manages per-layer K/V device tensors and position tracking. Native TensorRT +// KV engines update these buffers in place; legacy engines use present rows +// plus an attention mask. + +#include "runtime/models/smollm3/inference_state.h" +#include "trtmc/runtime/device_tensor.h" + +#include +#include +#include + +namespace trtmc { + +class ITrtModule; +using TrtModule = ITrtModule; + +// Explicit tensor names for KV cache I/O binding. +// Per-layer vectors hold expanded names; scalar names are for single inputs. +struct Smollm3KvCacheNames { + std::vector cache_k; + std::vector cache_v; + std::vector present_k; + std::vector present_v; + std::string cache_write_indices{"cache_write_indices"}; + std::string key_value_lengths{"key_value_lengths"}; + std::string position_id{"position_id"}; + std::string attention_mask{"attention_mask"}; +}; + +class Smollm3KvCache : public Smollm3InferenceState { + public: + // Allocate cache buffers for the given configuration. + // kv_dim = num_kv_heads * head_dim (size of one K or V row per layer). + // cache_dtype controls the element type for K/V cache buffers (default FP32). + // names provides explicit tensor names for engine I/O binding. + Smollm3KvCache(int32_t num_layers, int32_t max_length, int32_t kv_dim, cudaStream_t stream, + DType cache_dtype = DType::kFloat32, Smollm3KvCacheNames names = {}); + + // --- Smollm3InferenceState overrides --- + void reset() override; + void bind_to(TrtModule& module) override; + void prepare_step(TensorMap& inputs, int32_t seq_len = 1) override; + void advance(int32_t n_tokens = 1) override; + int32_t position() const override { return position_; } + int32_t max_length() const override { return max_length_; } + int32_t preferred_cache_rows() const override; + int32_t num_layers() const override { return num_layers_; } + bool needs_attention_mask() const override { return !native_kv_update_enabled_; } + std::size_t device_memory_bytes() const override; + const char* state_type() const override { return "dense_kv_cache"; } + bool ok() const override; + + // --- Smollm3KvCache-specific methods (not on the interface) --- + + // DEPRECATED: Use prepare_step() instead. + // Kept for backward compatibility with tests that call this directly. + void build_attention_mask(std::vector& mask) const; + + // Direct access for advanced use (cross-attention, VL embedding). + DeviceTensor& cache_k(int32_t layer) { return cache_k_[static_cast(layer)]; } + DeviceTensor& cache_v(int32_t layer) { return cache_v_[static_cast(layer)]; } + + // Complete a batched prefill and advance position_ to seq_len. Native + // TensorRT KV engines have already updated the aliased cache in place; + // legacy engines copy their per-layer outputs into the cache. + void write_prefill_kv(const std::vector& prefill_k, + const std::vector& prefill_v, int32_t seq_len); + + // Prepare a multi-token block whose new tokens may attend bidirectionally + // to one another while still seeing the valid prefix cache. + void prepare_bidirectional_step(TensorMap& inputs, int32_t seq_len); + + // Append batched present K/V at the current position. Used after causal + // block verification in diffusion-style text decoders. + void append_prefill_kv(const std::vector& prefill_k, + const std::vector& prefill_v, int32_t seq_len); + + // Move the logical cache length without touching device memory. Stale rows + // remain masked out by subsequent prepare_step calls. + void set_position(int32_t position); + + // Bind prefill KV tensors. Native TensorRT engines bind cache and present + // to the same full-capacity storage; legacy engines bind cache inputs only. + void bind_cache_inputs(TrtModule& module); + + private: + bool configure_binding_mode(TrtModule& module); + void validate_native_kv_contract(TrtModule& module) const; + void validate_native_aliases(const std::vector& present_k, + const std::vector& present_v) const; + void ensure_legacy_present_buffers(); + void bind_native_cache(TrtModule& module); + void write_native_kv_inputs(TensorMap& inputs, int32_t seq_len); + void rebind_cache_rows(int32_t cache_rows); + std::vector mask_shape_for_engine(int32_t mask_width, std::size_t mask_buf_size) const; + void write_position_input(TensorMap& inputs, int32_t seq_len); + void write_batched_mask(TensorMap& inputs, int32_t seq_len); + void write_bidirectional_mask(TensorMap& inputs, int32_t seq_len); + void write_decode_mask(TensorMap& inputs); + + std::vector cache_k_; // [num_layers], shape [max_length, kv_dim] + std::vector cache_v_; // [num_layers] + // Legacy-only single-step outputs, allocated lazily when a legacy engine is bound. + std::vector present_k_; + std::vector present_v_; + int32_t num_layers_{0}; + int32_t max_length_{0}; + int32_t kv_dim_{0}; + int32_t position_{0}; + cudaStream_t stream_{nullptr}; + // Buffers owned by this object — Tensor.data in prepare_step() points here. + std::vector mask_buf_; + std::vector pos_buf_vec_; + int32_t cache_write_index_{0}; + int32_t key_value_length_{0}; + bool has_position_input_{false}; + bool binding_mode_initialized_{false}; + bool native_kv_update_enabled_{false}; + bool dynamic_binding_enabled_{false}; + int32_t bound_cache_rows_{0}; + DType cache_dtype_{DType::kFloat32}; + std::size_t cache_element_size_{sizeof(float)}; + Smollm3KvCacheNames names_; + TrtModule* bound_module_{nullptr}; +}; + +} // namespace trtmc diff --git a/src/runtime/models/smollm3/pipeline.cpp b/src/runtime/models/smollm3/pipeline.cpp new file mode 100644 index 0000000000..1ab5bca6a3 --- /dev/null +++ b/src/runtime/models/smollm3/pipeline.cpp @@ -0,0 +1,1161 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/models/smollm3/pipeline.h" + +#include "runtime/models/smollm3/chat_templates.h" +#include "runtime/models/smollm3/kv_cache.h" +#include "runtime/models/smollm3/tensor_names.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace trtmc { + +namespace { + +struct StepTraceConfig { + bool enabled{false}; + std::string path; + int32_t start_position{0}; + int32_t end_position{std::numeric_limits::max()}; + int32_t top_k{8}; +}; + +// Process-wide step-trace state. Populated once from the resolved +// ConfigBundle by `apply_text_trace_config_from_registry` (below), called +// from the decoder plugin before pipeline construction. Replaces the +// TRTMC_TEXT_STEP_TRACE_* environment variables, which are now deleted. +StepTraceConfig& mutable_step_trace_config() { + static StepTraceConfig cfg; + return cfg; +} + +const StepTraceConfig& step_trace_config() { + return mutable_step_trace_config(); +} + +} // namespace + +// Called from decoder_plugin::create() with values resolved from +// ctx.runtime_config for the "text_trace" namespace. An empty path keeps +// tracing disabled. When a non-empty path is supplied, this truncates the +// target file so repeated runs don't concatenate. Not re-entrant; the +// caller serializes creation. +void apply_text_trace_config_from_registry(const std::string& path, int32_t start_position, + int32_t end_position, int32_t top_k) { + StepTraceConfig& cfg = mutable_step_trace_config(); + cfg.path = path; + cfg.enabled = !path.empty(); + cfg.start_position = start_position; + cfg.end_position = end_position; + cfg.top_k = std::max(int32_t{1}, top_k); + if (cfg.enabled) { + std::ofstream clear(cfg.path, std::ios::trunc); + } +} + +namespace { + +std::vector top_logit_indices(const std::vector& logits, int32_t top_n) { + std::vector order(logits.size()); + std::iota(order.begin(), order.end(), 0); + std::partial_sort( + order.begin(), order.begin() + top_n, order.end(), [&logits](int32_t lhs, int32_t rhs) { + if (logits[static_cast(lhs)] != logits[static_cast(rhs)]) { + return logits[static_cast(lhs)] > + logits[static_cast(rhs)]; + } + return lhs < rhs; + }); + return order; +} + +void write_step_trace_line(std::ostream& out, int32_t position_before, int32_t token_id, + int32_t decoder_idx, int32_t rows_before, int32_t rows_after, + const std::vector& logits, const std::vector& order, + int32_t top_n) { + out << "{\"position_before\":" << position_before << ",\"token_id\":" << token_id + << ",\"decoder_idx\":" << decoder_idx << ",\"rows_before\":" << rows_before + << ",\"rows_after\":" << rows_after << ",\"argmax_token\":" << order.front() + << ",\"argmax_logit\":" << logits[static_cast(order.front())] + << ",\"top_ids\":["; + for (int32_t i = 0; i < top_n; ++i) { + if (i > 0) + out << ','; + out << order[static_cast(i)]; + } + out << "],\"top_logits\":["; + for (int32_t i = 0; i < top_n; ++i) { + if (i > 0) + out << ','; + out << logits[static_cast(order[static_cast(i)])]; + } + out << "]}\n"; +} + +void maybe_append_step_trace(int32_t position_before, int32_t token_id, int32_t decoder_idx, + int32_t rows_before, int32_t rows_after, + const std::vector& logits) { + const auto& cfg = step_trace_config(); + if (!cfg.enabled || position_before < cfg.start_position || position_before > cfg.end_position) + return; + if (logits.empty()) + return; + const int32_t top_n = std::min(cfg.top_k, static_cast(logits.size())); + const auto order = top_logit_indices(logits, top_n); + std::ofstream out(cfg.path, std::ios::app); + if (!out) + return; + write_step_trace_line(out, position_before, token_id, decoder_idx, rows_before, rows_after, + logits, order, top_n); +} + +bool contains_boxed_answer(const std::string& text) { + const std::string marker = "\\boxed{"; + const auto start = text.find(marker); + if (start == std::string::npos) + return false; + return text.find('}', start + marker.size()) != std::string::npos; +} + +bool contains_final_answer(const std::string& text) { + const std::string marker = "Final answer:"; + const auto start = text.find(marker); + if (start == std::string::npos) + return false; + for (std::size_t i = start + marker.size(); i < text.size(); ++i) { + if (!std::isspace(static_cast(text[i]))) + return true; + } + return false; +} + +std::vector +single_decoder_context(std::unique_ptr decoder) { + std::vector decoders; + decoders.push_back(Smollm3TextGenerationPipeline::DecoderContext{0, std::move(decoder)}); + return decoders; +} + +Smollm3TextGenConfig normalize_eos_token_ids(Smollm3TextGenConfig config) { + if (config.id_eos_ids.empty() && config.id_eos >= 0) + config.id_eos_ids.push_back(config.id_eos); + if (!config.id_eos_ids.empty()) + config.id_eos = config.id_eos_ids.front(); + return config; +} + +std::string normalize_generation_mode(std::string mode) { + std::transform(mode.begin(), mode.end(), mode.begin(), + [](unsigned char ch) { return static_cast(std::tolower(ch)); }); + std::replace(mode.begin(), mode.end(), '-', '_'); + return mode; +} + +bool greedy_text_diffusion_params(const Smollm3SamplingParams& params) { + return params.seed < 0 && + (params.temperature <= 1e-6F || + (params.top_k <= 1 && params.top_p >= 1.0F - 1e-6F && params.min_p <= 1e-6F)); +} + +struct TokenConfidence { + int32_t pos{0}; + int32_t token_id{0}; + float confidence{0.0F}; +}; + +TokenConfidence argmax_with_confidence(const float* logits, int32_t vocab, int32_t pos) { + TokenConfidence out; + out.pos = pos; + if (logits == nullptr || vocab <= 0) + return out; + int32_t best = 0; + float max_logit = logits[0]; + for (int32_t i = 1; i < vocab; ++i) { + if (logits[i] > max_logit) { + max_logit = logits[i]; + best = i; + } + } + double denom = 0.0; + for (int32_t i = 0; i < vocab; ++i) + denom += std::exp(static_cast(logits[i] - max_logit)); + out.token_id = best; + out.confidence = denom > 0.0 ? static_cast(1.0 / denom) : 0.0F; + return out; +} + +std::vector transfer_quota_schedule(int32_t masked, int32_t steps) { + steps = std::max(steps, 1); + std::vector quota(static_cast(steps), 0); + const int32_t base = masked / steps; + const int32_t rem = masked % steps; + for (int32_t i = 0; i < steps; ++i) + quota[static_cast(i)] = base + (i < rem ? 1 : 0); + return quota; +} + +std::vector masked_predictions(const std::vector& logits, + const std::vector& block, + int32_t mask_token_id, int32_t vocab_size) { + std::vector preds; + if (vocab_size <= 0) + return preds; + const auto rows = static_cast(logits.size() / static_cast(vocab_size)); + const int32_t usable = std::min(rows, static_cast(block.size())); + preds.reserve(static_cast(usable)); + for (int32_t i = 0; i < usable; ++i) { + if (block[static_cast(i)] != mask_token_id) + continue; + preds.push_back(argmax_with_confidence( + logits.data() + static_cast(i) * static_cast(vocab_size), + vocab_size, i)); + } + std::sort(preds.begin(), preds.end(), + [](const TokenConfidence& lhs, const TokenConfidence& rhs) { + if (lhs.confidence != rhs.confidence) + return lhs.confidence > rhs.confidence; + return lhs.pos < rhs.pos; + }); + return preds; +} + +void apply_diffusion_transfer(std::vector& block, + const std::vector& preds, int32_t quota, + bool use_threshold, float threshold) { + if (preds.empty()) + return; + if (use_threshold) { + block[static_cast(preds.front().pos)] = preds.front().token_id; + for (std::size_t i = 1; i < preds.size(); ++i) { + if (preds[i].confidence >= threshold) + block[static_cast(preds[i].pos)] = preds[i].token_id; + } + return; + } + quota = std::max(0, std::min(quota, static_cast(preds.size()))); + for (int32_t i = 0; i < quota; ++i) + block[static_cast(preds[static_cast(i)].pos)] = + preds[static_cast(i)].token_id; +} + +void apply_linear_spec_transfer(std::vector& block, + const std::vector& preds, bool threshold_enabled, + float threshold) { + if (preds.empty()) + return; + if (!threshold_enabled) { + for (const auto& pred : preds) + block[static_cast(pred.pos)] = pred.token_id; + return; + } + + bool changed = false; + for (const auto& pred : preds) { + if (pred.confidence >= threshold) { + block[static_cast(pred.pos)] = pred.token_id; + changed = true; + } + } + if (!changed) + block[static_cast(preds.front().pos)] = preds.front().token_id; +} + +bool has_mask_token(const std::vector& block, int32_t mask_token_id) { + return std::find(block.begin(), block.end(), mask_token_id) != block.end(); +} + +} // namespace + +Smollm3TextGenerationPipeline::Smollm3TextGenerationPipeline( + std::unique_ptr decoder, std::unique_ptr state, + Smollm3TextGenConfig config, cudaStream_t stream, std::shared_ptr tokenizer, + std::string model_id_str, std::unique_ptr sampler, + std::shared_ptr distributed_owner) + : Smollm3TextGenerationPipeline(single_decoder_context(std::move(decoder)), std::move(state), + std::move(config), stream, std::move(tokenizer), + std::move(model_id_str), std::move(sampler), + /*prefill=*/nullptr, /*linear_spec_lora_prefill=*/nullptr, + std::move(distributed_owner)) {} + +Smollm3TextGenerationPipeline::Smollm3TextGenerationPipeline( + std::vector decoders, std::unique_ptr state, + Smollm3TextGenConfig config, cudaStream_t stream, std::shared_ptr tokenizer, + std::string model_id_str, std::unique_ptr sampler, + std::unique_ptr prefill, std::unique_ptr linear_spec_lora_prefill, + std::shared_ptr distributed_owner) + : distributed_owner_(std::move(distributed_owner)), decoders_(std::move(decoders)), + prefill_(std::move(prefill)), linear_spec_lora_prefill_(std::move(linear_spec_lora_prefill)), + state_(std::move(state)), config_(normalize_eos_token_ids(std::move(config))), + stream_(stream), tokenizer_(std::move(tokenizer)), model_id_(std::move(model_id_str)), + sampler_(std::move(sampler)), logits_output_name_(config_.logits_output_name) { + if (decoders_.empty()) { + throw std::runtime_error("Smollm3TextGenerationPipeline: no decoder modules"); + } + for (const auto& decoder_ctx : decoders_) { + if (!decoder_ctx.module || !decoder_ctx.module->ok()) { + throw std::runtime_error("Smollm3TextGenerationPipeline: invalid decoder module"); + } + } + if (!state_ || !state_->ok()) { + throw std::runtime_error("Smollm3TextGenerationPipeline: invalid inference state"); + } + + // CUDA Graphs: capture TRT kernels on first step, replay on subsequent + // steps. Disabled via --set runtime.disable_cuda_graph=true (replaces + // the deleted TRTMC_DISABLE_CUDA_GRAPH env var). + if (!config_.disable_cuda_graph) { + for (auto& decoder_ctx : decoders_) + decoder_ctx.module->enable_cuda_graph(); + } + + // GPU-side argmax is only valid for truly greedy decoding. Populated + // from runtime.prefer_gpu_greedy (replaces the deleted TRTMC_GPU_ARGMAX + // env var). We record the preference here and instantiate per-call + // when the requested sampling parameters are actually greedy. + prefer_gpu_greedy_ = config_.prefer_gpu_greedy; +} + +// Encode a prompt, optionally applying a chat template first. +// Deduplicates the leading BOS token that chat templates embed but +// the tokenizer's add_special_tokens may also prepend. +static std::vector encode_prompt(const ITokenizer& tokenizer, + const Smollm3TextGenConfig& config, + const std::string& prompt, const GenerateConfig& cfg) { + std::string effective = prompt; + bool templated = false; + if (cfg.use_chat_template && !config.chat_template_format.empty()) { + effective = + smollm3_apply_chat_template(config.chat_template_format, prompt, cfg.enable_thinking); + templated = true; + } + auto ids = tokenizer.encode(effective); + if (templated && ids.size() >= 2 && config.id_bos >= 0 && ids[0] == config.id_bos && + ids[1] == config.id_bos) { + ids.erase(ids.begin()); + } + return ids; +} + +TextResult Smollm3TextGenerationPipeline::generate(const std::string& prompt, + const GenerateConfig& cfg) { + if (!tokenizer_) { + throw std::runtime_error("Smollm3TextGenerationPipeline: no tokenizer configured"); + } + + auto input_ids = encode_prompt(*tokenizer_, config_, prompt, cfg); + int32_t max_new = (cfg.max_new_tokens > 0) ? cfg.max_new_tokens : 128; + auto sp = smollm3_sampling_params_from_config(cfg, config_.id_eos_ids); + last_setup_ms_ = 0.0; + auto timed = generate_from_ids(input_ids, max_new, sp, cfg); + + // Decode only the NEW tokens (skip input) + std::vector new_tokens(timed.token_ids.begin() + + static_cast(input_ids.size()), + timed.token_ids.end()); + std::string text = tokenizer_->decode(new_tokens); + + auto result = + TextResult{std::move(text), std::move(new_tokens), timed.prefill_ms, timed.decode_ms}; + result.setup_ms = last_setup_ms_; + return result; +} + +Smollm3TextGenerationPipeline::GenerationResult +Smollm3TextGenerationPipeline::generate_ids(const std::vector& input_ids, + const GenerateConfig& cfg) { + int32_t max_new = cfg.max_new_tokens; // honour exact value (0 = no generation) + auto sp = smollm3_sampling_params_from_config(cfg, config_.id_eos_ids); + return GenerationResult{generate_from_ids(input_ids, max_new, sp, cfg).token_ids}; +} + +std::unique_ptr +Smollm3TextGenerationPipeline::make_step_sampler(const Smollm3SamplingParams& params) { + const bool greedy_params = + (params.temperature < 1e-6F) || + (params.top_k <= 1 && params.top_p >= 1.0F && params.min_p <= 0.0F && params.seed < 0); + if (prefer_gpu_greedy_ && greedy_params) { + if (auto gpu = create_smollm3_gpu_greedy_sampler(stream_)) + return gpu; + } + return create_smollm3_sampler(params); +} + +// Helper: gather per-layer present_k/present_v device pointers from the +// prefill TrtModule. Returns false if any layer's tensor is missing — in +// that case the caller falls back to the per-token decode loop. +namespace { +bool gather_prefill_kv_pointers(TrtModule& prefill, const Smollm3TextGenConfig& cfg, + std::vector& pk, std::vector& pv) { + pk.resize(static_cast(cfg.num_layers)); + pv.resize(static_cast(cfg.num_layers)); + for (int32_t i = 0; i < cfg.num_layers; ++i) { + const auto li = static_cast(i); + pk[li] = prefill.device_ptr(smollm3_expand_layer_name(cfg.present_k_pattern, i)); + pv[li] = prefill.device_ptr(smollm3_expand_layer_name(cfg.present_v_pattern, i)); + if (pk[li] == nullptr || pv[li] == nullptr) + return false; + } + return true; +} + +bool batched_prefill_supported(const TrtModule* prefill, const Smollm3TextGenConfig& cfg, + int32_t sq, Smollm3InferenceState* state) { + if (prefill == nullptr || sq <= 0) + return false; + if (cfg.num_layers <= 0 || cfg.vocab_size <= 0) + return false; + return dynamic_cast(state) != nullptr; +} + +int32_t resolve_prefill_chunk_limit(const Smollm3KvCache& kv, const Smollm3TextGenConfig& cfg, + int32_t sq) { + if (kv.needs_attention_mask()) { + if (cfg.prefill_max_length > 0 && sq > cfg.prefill_max_length) + return 0; + return sq; + } + if (cfg.prefill_max_length <= 0) + throw std::runtime_error("SmolLM3 native KV prefill engine has no valid profile capacity"); + return cfg.prefill_max_length; +} + +void validate_generation_capacity(const std::vector& input_ids, int32_t max_new_tokens, + Smollm3InferenceState* state, const TrtModule* module) { + if (module == nullptr || !module->has_input("cache_write_indices") || + !module->has_input("key_value_lengths")) { + return; + } + const auto* kv = dynamic_cast(state); + if (kv == nullptr) + return; + + const auto capacity = static_cast(kv->max_length()); + if (input_ids.size() > capacity || + (max_new_tokens > 0 && + static_cast(max_new_tokens) > capacity - input_ids.size())) { + throw std::runtime_error( + "SmolLM3 requested prompt and generation exceed the model's fixed KV cache capacity"); + } +} +} // namespace + +bool Smollm3TextGenerationPipeline::run_prefill_batched(const std::vector& input_ids, + std::vector& logits, + bool retain_device_logits) { + const auto sq = static_cast(input_ids.size()); + if (!batched_prefill_supported(prefill_.get(), config_, sq, state_.get())) + return false; + auto* kv = static_cast(state_.get()); + + // The prefill module shares the same external KV cache buffers as the + // decode module(s), so we rebind the cache_k/cache_v inputs onto the + // prefill execution context before running. + kv->bind_cache_inputs(*prefill_); + if (sq > kv->max_length()) { + throw std::runtime_error("SmolLM3 sequence exceeds the model's fixed KV cache capacity"); + } + + std::vector pk, pv; + if (!gather_prefill_kv_pointers(*prefill_, config_, pk, pv)) + return false; + + const int32_t chunk_limit = resolve_prefill_chunk_limit(*kv, config_, sq); + if (chunk_limit == 0) + return false; + + int32_t chunk_count = 0; + int32_t max_chunk_size = 0; + for (int32_t start = 0; start < sq;) { + const int32_t chunk_size = std::min(chunk_limit, sq - start); + run_prefill_chunk(input_ids.data() + start, chunk_size, pk, pv, *kv, logits, + retain_device_logits); + ++chunk_count; + max_chunk_size = std::max(max_chunk_size, chunk_size); + start += chunk_size; + } + + log_batched_prefill(sq, chunk_count, max_chunk_size); + return true; +} + +void Smollm3TextGenerationPipeline::run_prefill_chunk(const int32_t* token_ids, int32_t chunk_size, + const std::vector& present_k, + const std::vector& present_v, + Smollm3KvCache& kv, + std::vector& logits, + bool retain_device_logits) { + TensorMap inputs; + Tensor token_tensor; + token_tensor.data = const_cast(token_ids); + token_tensor.shape = {static_cast(chunk_size)}; + token_tensor.dtype = DType::kInt32; + inputs[config_.token_id_name] = token_tensor; + state_->prepare_step(inputs, chunk_size); + + TensorMap outputs = prefill_->forward(inputs); + const auto logits_it = outputs.find(config_.logits_output_name); + if (logits_it == outputs.end()) { + throw std::runtime_error( + "Smollm3TextGenerationPipeline: prefill module has no logits output"); + } + + const auto& logits_tensor = logits_it->second; + const auto vocab = static_cast(config_.vocab_size); + if (static_cast(logits_tensor.numel()) < vocab) { + throw std::runtime_error( + "Smollm3TextGenerationPipeline: prefill logits are smaller than vocabulary"); + } + + logits.resize(vocab); + const auto logits_offset = static_cast(logits_tensor.numel()) - vocab; + std::memcpy(logits.data(), static_cast(logits_tensor.data) + logits_offset, + vocab * sizeof(float)); + if (retain_device_logits) { + const auto* device_logits = + static_cast(prefill_->device_ptr(config_.logits_output_name)); + if (device_logits == nullptr) { + throw std::runtime_error( + "Smollm3TextGenerationPipeline: prefill logits have no device buffer"); + } + d_logits_ptr_ = device_logits + logits_offset; + } + kv.append_prefill_kv(present_k, present_v, chunk_size); +} + +void Smollm3TextGenerationPipeline::log_batched_prefill(int32_t token_count, int32_t chunk_count, + int32_t max_chunk_size) const { + std::cerr << "[trtmc.prefill] tokens=" << token_count << " launches=" << chunk_count + << " max_chunk=" << max_chunk_size << '\n'; + if (!config_.log_runtime_stats) + return; + + std::cerr << "[trtmc] Batched prefill ("; + if (!config_.prefill_log_label.empty()) { + std::cerr << config_.prefill_log_label; + } else { + std::cerr << "profile " << config_.prefill_profile_index; + } + std::cerr << "): " << token_count << " tokens in " << chunk_count << " call"; + if (chunk_count != 1) + std::cerr << 's'; + std::cerr << " (max chunk=" << max_chunk_size << ")\n"; +} + +const TrtModule* Smollm3TextGenerationPipeline::generation_capacity_module() const { + if (prefill_ != nullptr) + return prefill_.get(); + return decoders_.front().module.get(); +} + +void Smollm3TextGenerationPipeline::prime_decoder_after_batched_prefill( + const std::vector& input_ids) { + if (input_ids.empty()) + return; + + TrtModule& decoder = bind_decoder_for_step(); + if (!decoder.cuda_graph_active()) + return; + + int32_t token_id = input_ids.back(); + TensorMap inputs; + Tensor token_tensor; + token_tensor.data = &token_id; + token_tensor.shape = {1}; + token_tensor.dtype = DType::kInt32; + inputs[config_.token_id_name] = token_tensor; + + state_->prepare_step(inputs); + decoder.forward_async(inputs); + decoder.sync(); +} + +void Smollm3TextGenerationPipeline::run_prefill(const std::vector& input_ids, + std::vector& logits, bool gpu_sampling) { + // Fast path: batched prefill writes K/V in profile-bounded chunks and + // exposes last-token logits on the sampler's requested host or device path. + if (run_prefill_batched(input_ids, logits, gpu_sampling)) { + prime_decoder_after_batched_prefill(input_ids); + state_->mark_prefill_complete(); + return; + } + for (std::size_t i = 0; i + 1 < input_ids.size(); ++i) { + if (gpu_sampling) + run_step_device(input_ids[i]); + else + run_step(input_ids[i], logits); + } + const int32_t last_token = input_ids.back(); + if (gpu_sampling) + run_step_device(last_token); + else + run_step(last_token, logits); + state_->mark_prefill_complete(); +} + +TrtModule& Smollm3TextGenerationPipeline::require_block_prefill(int32_t sq, + TrtModule* prefill_override) { + TrtModule* prefill = prefill_override != nullptr ? prefill_override : prefill_.get(); + if (prefill == nullptr) + throw std::runtime_error( + "Smollm3TextGenerationPipeline: block generation requires prefill module"); + if (sq <= 0) + throw std::runtime_error("Smollm3TextGenerationPipeline: empty block"); + if (config_.prefill_max_length > 0 && sq > config_.prefill_max_length) { + throw std::runtime_error( + "Smollm3TextGenerationPipeline: block length exceeds prefill profile"); + } + return *prefill; +} + +Smollm3KvCache& Smollm3TextGenerationPipeline::require_block_kv_cache() { + auto* kv = dynamic_cast(state_.get()); + if (kv == nullptr) + throw std::runtime_error( + "Smollm3TextGenerationPipeline: block generation requires Smollm3KvCache"); + return *kv; +} + +void Smollm3TextGenerationPipeline::copy_block_logits(const TensorMap& outputs, + std::vector& logits) const { + auto logits_it = outputs.find(config_.logits_output_name); + if (logits_it == outputs.end()) + throw std::runtime_error("Smollm3TextGenerationPipeline: prefill module has no '" + + config_.logits_output_name + "' output"); + + const auto& lt = logits_it->second; + const auto num_logits = static_cast(lt.numel()); + logits.resize(num_logits); + std::memcpy(logits.data(), lt.data, num_logits * sizeof(float)); +} + +void Smollm3TextGenerationPipeline::append_prefill_kv(Smollm3KvCache& kv, TrtModule& prefill, + int32_t sq) { + std::vector pk, pv; + if (!gather_prefill_kv_pointers(prefill, config_, pk, pv)) { + throw std::runtime_error( + "Smollm3TextGenerationPipeline: prefill module is missing present_k/present_v outputs"); + } + kv.append_prefill_kv(pk, pv, sq); +} + +void Smollm3TextGenerationPipeline::run_prefill_block(const std::vector& input_ids, + bool bidirectional, bool append_kv, + std::vector& logits, + TrtModule* prefill_override) { + const auto sq = static_cast(input_ids.size()); + TrtModule& prefill = require_block_prefill(sq, prefill_override); + Smollm3KvCache& kv = require_block_kv_cache(); + + kv.bind_cache_inputs(prefill); + + TensorMap inputs; + Tensor tok_t; + tok_t.data = const_cast(input_ids.data()); + tok_t.shape = {static_cast(sq)}; + tok_t.dtype = DType::kInt32; + inputs[config_.token_id_name] = tok_t; + if (bidirectional) + kv.prepare_bidirectional_step(inputs, sq); + else + kv.prepare_step(inputs, sq); + + copy_block_logits(prefill.forward(inputs), logits); + if (append_kv) + append_prefill_kv(kv, prefill, sq); +} + +std::string +Smollm3TextGenerationPipeline::resolve_generation_mode(const GenerateConfig& cfg) const { + std::string mode = normalize_generation_mode(cfg.text_generation_mode); + if (mode.empty()) + mode = "auto"; + if (mode == "auto" && config_.supports_text_diffusion) + mode = "diffusion"; + if (mode == "autoregressive") + mode = "ar"; + if (mode == "linear_speculation") + mode = "linear_spec"; + if (mode == "linear_speculation_lora" || mode == "linear_spec_adapter") + mode = "linear_spec_lora"; + return mode; +} + +void Smollm3TextGenerationPipeline::reset_generation_context() { + using Clock = std::chrono::steady_clock; + const auto start = Clock::now(); + state_->reset(); + d_logits_ptr_ = nullptr; + state_bound_ = false; + for (auto& decoder_ctx : decoders_) + decoder_ctx.module->reset_execution_context(); + if (prefill_) + prefill_->reset_execution_context(); + if (linear_spec_lora_prefill_) + linear_spec_lora_prefill_->reset_execution_context(); + last_setup_ms_ = std::chrono::duration(Clock::now() - start).count(); +} + +int32_t Smollm3TextGenerationPipeline::resolve_text_diffusion_block_length( + const GenerateConfig& cfg, int32_t max_new_tokens, bool require_divisible) const { + if (!config_.supports_text_diffusion || config_.mask_token_id < 0) + throw std::runtime_error( + "Smollm3TextGenerationPipeline: bundle does not support text diffusion"); + const int32_t block_len = + cfg.block_length > 0 ? cfg.block_length : std::max(config_.diffusion_block_length, 1); + if (require_divisible && max_new_tokens % block_len != 0) { + throw std::runtime_error("Smollm3TextGenerationPipeline: diffusion mode requires " + "max_new_tokens % block_length == 0"); + } + return block_len; +} + +int32_t Smollm3TextGenerationPipeline::seed_next_token_from_prefill( + const std::vector& input_ids, std::vector& logits, int32_t vocab) { + run_prefill_block(input_ids, /*bidirectional=*/false, /*append_kv=*/true, logits); + if (static_cast(logits.size()) < vocab) + throw std::runtime_error("Smollm3TextGenerationPipeline: missing prefill logits"); + return argmax_with_confidence(logits.data() + logits.size() - static_cast(vocab), + vocab, 0) + .token_id; +} + +void Smollm3TextGenerationPipeline::fill_diffusion_block(std::vector& block, + std::vector& logits, + int32_t block_len, int32_t vocab, + bool use_threshold, float threshold) { + const int32_t initial_masked = block_len - 1; + const auto quotas = transfer_quota_schedule(initial_masked, block_len); + for (int32_t step = 0; step < block_len && has_mask_token(block, config_.mask_token_id); + ++step) { + run_prefill_block(block, /*bidirectional=*/true, /*append_kv=*/false, logits); + if (static_cast(logits.size()) < block_len * vocab) { + throw std::runtime_error( + "Smollm3TextGenerationPipeline: diffusion engine must output full block logits"); + } + const auto preds = masked_predictions(logits, block, config_.mask_token_id, vocab); + apply_diffusion_transfer(block, preds, quotas[static_cast(step)], + use_threshold, threshold); + } +} + +int32_t Smollm3TextGenerationPipeline::verify_diffusion_block(const std::vector& block, + std::vector& logits, + int32_t block_len, int32_t vocab) { + run_prefill_block(block, /*bidirectional=*/false, /*append_kv=*/true, logits); + if (static_cast(logits.size()) < block_len * vocab) { + throw std::runtime_error( + "Smollm3TextGenerationPipeline: diffusion engine must output full verify logits"); + } + return argmax_with_confidence(logits.data() + (static_cast(block_len - 1) * + static_cast(vocab)), + vocab, block_len - 1) + .token_id; +} + +bool Smollm3TextGenerationPipeline::append_tokens_until_eos( + const std::vector& tokens, std::vector& output, + const Smollm3SamplingParams& params) const { + for (int32_t token : tokens) { + output.push_back(token); + if (smollm3_is_eos_token(params, token)) + return true; + } + return false; +} + +void Smollm3TextGenerationPipeline::fill_linear_spec_block(std::vector& block, + std::vector& logits, + int32_t block_len, int32_t vocab, + bool threshold_enabled, float threshold, + bool use_lora_draft) { + while (has_mask_token(block, config_.mask_token_id)) { + TrtModule* draft_prefill = use_lora_draft ? linear_spec_lora_prefill_.get() : nullptr; + run_prefill_block(block, /*bidirectional=*/true, /*append_kv=*/false, logits, + draft_prefill); + if (static_cast(logits.size()) < block_len * vocab) { + throw std::runtime_error( + "Smollm3TextGenerationPipeline: linear_spec engine must output full block logits"); + } + const auto preds = masked_predictions(logits, block, config_.mask_token_id, vocab); + apply_linear_spec_transfer(block, preds, threshold_enabled, threshold); + } +} + +std::vector +Smollm3TextGenerationPipeline::verify_linear_spec_block(const std::vector& block, + std::vector& logits, + int32_t block_len, int32_t vocab) { + run_prefill_block(block, /*bidirectional=*/false, /*append_kv=*/true, logits); + if (static_cast(logits.size()) < block_len * vocab) { + throw std::runtime_error( + "Smollm3TextGenerationPipeline: linear_spec engine must output full verify logits"); + } + + std::vector ar_tokens; + ar_tokens.reserve(static_cast(block_len)); + for (int32_t i = 0; i < block_len; ++i) { + ar_tokens.push_back( + argmax_with_confidence( + logits.data() + (static_cast(i) * static_cast(vocab)), + vocab, i) + .token_id); + } + return ar_tokens; +} + +int32_t +Smollm3TextGenerationPipeline::count_linear_spec_accepts(const std::vector& ar_tokens, + const std::vector& block) { + if (ar_tokens.empty()) + return 0; + if (block.size() < 2) + return 1; + int32_t accepted = 0; + const auto limit = static_cast(std::min(ar_tokens.size(), block.size() - 1)); + for (int32_t i = 0; i < limit; ++i) { + if (ar_tokens[static_cast(i)] != block[static_cast(i + 1)]) + break; + ++accepted; + } + return accepted + 1; +} + +bool Smollm3TextGenerationPipeline::append_linear_spec_tokens( + const std::vector& ar_tokens, int32_t emit_count, std::vector& output, + int32_t& generated, const Smollm3SamplingParams& params) const { + for (int32_t i = 0; i < emit_count; ++i) { + const int32_t token = ar_tokens[static_cast(i)]; + output.push_back(token); + ++generated; + if (smollm3_is_eos_token(params, token)) + return true; + } + return false; +} + +Smollm3TextGenerationPipeline::TimedGenResult Smollm3TextGenerationPipeline::generate_from_ids( + const std::vector& input_ids, int32_t max_new_tokens, + const Smollm3SamplingParams& params, const GenerateConfig& cfg) { + using Clock = std::chrono::steady_clock; + if (max_new_tokens == 0 || input_ids.empty()) + return TimedGenResult{input_ids, 0.0, 0.0}; + validate_generation_capacity(input_ids, max_new_tokens, state_.get(), + generation_capacity_module()); + + const std::string mode = resolve_generation_mode(cfg); + if (mode == "diffusion" || mode == "dlm") + return generate_diffusion_from_ids(input_ids, max_new_tokens, params, cfg); + if (mode == "linear_spec" || mode == "linear_spec_lora") + return generate_linear_spec_from_ids(input_ids, max_new_tokens, params, cfg, + mode == "linear_spec_lora"); + if (mode != "auto" && mode != "ar") + throw std::runtime_error("Smollm3TextGenerationPipeline: unsupported generation mode '" + + mode + "'"); + + Smollm3ISampler* active_sampler = sampler_.get(); + std::unique_ptr local_sampler; + if (!active_sampler) { + local_sampler = make_step_sampler(params); + active_sampler = local_sampler.get(); + } + active_sampler->reset(); + + reset_generation_context(); + state_->set_prompt_length(static_cast(input_ids.size())); + + std::vector logits; + const bool gpu_sampling = (active_sampler->logits_location() == Smollm3LogitsLocation::DEVICE); + const auto t0 = Clock::now(); + run_prefill(input_ids, logits, gpu_sampling); + const auto t1 = Clock::now(); + + std::vector output = input_ids; + run_decode_loop(active_sampler, params, output, logits, max_new_tokens, gpu_sampling, cfg, + static_cast(input_ids.size())); + const auto t2 = Clock::now(); + + const double prefill_ms = std::chrono::duration(t1 - t0).count(); + const double decode_ms = std::chrono::duration(t2 - t1).count(); + return TimedGenResult{std::move(output), prefill_ms, decode_ms}; +} + +Smollm3TextGenerationPipeline::TimedGenResult +Smollm3TextGenerationPipeline::generate_diffusion_from_ids(const std::vector& input_ids, + int32_t max_new_tokens, + const Smollm3SamplingParams& params, + const GenerateConfig& cfg) { + using Clock = std::chrono::steady_clock; + if (!greedy_text_diffusion_params(params)) { + throw std::runtime_error( + "Smollm3TextGenerationPipeline: diffusion mode currently supports greedy temperature=0 " + "generation"); + } + const int32_t block_len = + resolve_text_diffusion_block_length(cfg, max_new_tokens, /*require_divisible=*/true); + const bool use_threshold = cfg.confidence_threshold >= 0.0F; + const float threshold = cfg.confidence_threshold; + const int32_t vocab = config_.vocab_size; + + reset_generation_context(); + state_->set_prompt_length(static_cast(input_ids.size())); + + std::vector logits; + const auto t0 = Clock::now(); + int32_t next_token = seed_next_token_from_prefill(input_ids, logits, vocab); + const auto t1 = Clock::now(); + + std::vector output = input_ids; + const int32_t num_blocks = max_new_tokens / block_len; + const auto decode_start = Clock::now(); + for (int32_t block_idx = 0; block_idx < num_blocks; ++block_idx) { + std::vector block(static_cast(block_len), config_.mask_token_id); + block[0] = next_token; + fill_diffusion_block(block, logits, block_len, vocab, use_threshold, threshold); + next_token = verify_diffusion_block(block, logits, block_len, vocab); + + if (append_tokens_until_eos(block, output, params)) { + const auto t2 = Clock::now(); + return TimedGenResult{ + std::move(output), std::chrono::duration(t1 - t0).count(), + std::chrono::duration(t2 - decode_start).count()}; + } + } + + const auto t2 = Clock::now(); + return TimedGenResult{std::move(output), + std::chrono::duration(t1 - t0).count(), + std::chrono::duration(t2 - decode_start).count()}; +} + +Smollm3TextGenerationPipeline::TimedGenResult +Smollm3TextGenerationPipeline::generate_linear_spec_from_ids(const std::vector& input_ids, + int32_t max_new_tokens, + const Smollm3SamplingParams& params, + const GenerateConfig& cfg, + bool use_lora_draft) { + using Clock = std::chrono::steady_clock; + if (!greedy_text_diffusion_params(params)) { + throw std::runtime_error("Smollm3TextGenerationPipeline: linear_spec mode currently " + "supports greedy temperature=0 " + "generation"); + } + if (use_lora_draft && linear_spec_lora_prefill_ == nullptr) { + throw std::runtime_error("Smollm3TextGenerationPipeline: linear_spec_lora mode requires a " + "linear-spec LoRA engine"); + } + const int32_t block_len = + resolve_text_diffusion_block_length(cfg, max_new_tokens, /*require_divisible=*/false); + const bool threshold_enabled = cfg.confidence_threshold > 0.0F; + const float threshold = cfg.confidence_threshold; + const int32_t vocab = config_.vocab_size; + + reset_generation_context(); + state_->set_prompt_length(static_cast(input_ids.size())); + + std::vector logits; + const auto t0 = Clock::now(); + int32_t next_token = seed_next_token_from_prefill(input_ids, logits, vocab); + const auto t1 = Clock::now(); + + std::vector output = input_ids; + output.push_back(next_token); + if (smollm3_is_eos_token(params, next_token)) { + return TimedGenResult{std::move(output), + std::chrono::duration(t1 - t0).count(), 0.0}; + } + + auto* kv = dynamic_cast(state_.get()); + if (kv == nullptr) + throw std::runtime_error( + "Smollm3TextGenerationPipeline: linear_spec requires Smollm3KvCache"); + + int32_t generated = 1; + const auto decode_start = Clock::now(); + while (generated < max_new_tokens) { + const int32_t cache_len = kv->position(); + std::vector block(static_cast(block_len), config_.mask_token_id); + block[0] = next_token; + + fill_linear_spec_block(block, logits, block_len, vocab, threshold_enabled, threshold, + use_lora_draft); + const auto ar_tokens = verify_linear_spec_block(block, logits, block_len, vocab); + const int32_t accepted = count_linear_spec_accepts(ar_tokens, block); + const int32_t emit_count = std::min(accepted, max_new_tokens - generated); + kv->set_position(cache_len + emit_count); + next_token = ar_tokens[static_cast(emit_count - 1)]; + + if (append_linear_spec_tokens(ar_tokens, emit_count, output, generated, params)) { + const auto t2 = Clock::now(); + return TimedGenResult{ + std::move(output), std::chrono::duration(t1 - t0).count(), + std::chrono::duration(t2 - decode_start).count()}; + } + } + + const auto t2 = Clock::now(); + return TimedGenResult{std::move(output), + std::chrono::duration(t1 - t0).count(), + std::chrono::duration(t2 - decode_start).count()}; +} + +bool Smollm3TextGenerationPipeline::should_stop_on_answer(const std::vector& output, + int32_t prompt_token_count, + const GenerateConfig& cfg, int32_t steps, + int32_t stop_interval, + bool is_eos) const { + if (!cfg.stop_on_boxed_answer || !tokenizer_) + return false; + if ((steps % stop_interval) != 0 && !is_eos) + return false; + std::vector new_tokens(output.begin() + prompt_token_count, output.end()); + const std::string decoded = tokenizer_->decode(new_tokens); + return contains_boxed_answer(decoded) || contains_final_answer(decoded); +} + +void Smollm3TextGenerationPipeline::log_decode_summary(int32_t steps, double ms) const { + if (steps <= 0 || !config_.log_runtime_stats) + return; + const double tps = steps * 1000.0 / ms; + const bool cuda_graph_on = + active_decoder_index_ >= 0 && + decoders_[static_cast(active_decoder_index_)].module->cuda_graph_active(); + std::cerr << "[trtmc] Decode: " << steps << " tokens, " << ms << " ms, " << tps << " tok/s" + << (cuda_graph_on ? " [CUDA Graph ON]" : "") << '\n'; +} + +int32_t Smollm3TextGenerationPipeline::run_decode_loop( + Smollm3ISampler* sampler, const Smollm3SamplingParams& params, std::vector& output, + std::vector& logits, int32_t max_new_tokens, bool gpu_sampling, + const GenerateConfig& cfg, int32_t prompt_token_count) { + const int32_t vocab_size = + gpu_sampling ? config_.vocab_size : static_cast(logits.size()); + const int32_t stop_interval = std::max(cfg.stop_check_interval, 1); + const auto decode_start = std::chrono::steady_clock::now(); + int32_t steps = 0; + for (int32_t step = 0; step < max_new_tokens; ++step) { + const float* sample_ptr = gpu_sampling ? d_logits_ptr_ : logits.data(); + const Smollm3SampleResult result = sampler->sample(sample_ptr, vocab_size, params); + const bool is_eos = result.is_eos || smollm3_is_eos_token(params, result.token_id); + output.push_back(result.token_id); + ++steps; + if (should_stop_on_answer(output, prompt_token_count, cfg, steps, stop_interval, is_eos)) + break; + if (is_eos) + break; + if (gpu_sampling) + run_step_device(result.token_id); + else + run_step(result.token_id, logits); + } + const auto decode_end = std::chrono::steady_clock::now(); + const double ms = std::chrono::duration(decode_end - decode_start).count(); + log_decode_summary(steps, ms); + return steps; +} + +int32_t Smollm3TextGenerationPipeline::select_decoder_index(int32_t desired_rows) const { + if (decoders_.size() == 1) + return 0; + + int32_t fallback_idx = 0; + int32_t fallback_rows = std::numeric_limits::max(); + for (std::size_t i = 0; i < decoders_.size(); ++i) { + const int32_t kv_rows = decoders_[i].kv_rows; + if (kv_rows == desired_rows) + return static_cast(i); + if (kv_rows > 0 && kv_rows >= desired_rows && kv_rows < fallback_rows) { + fallback_rows = kv_rows; + fallback_idx = static_cast(i); + } + } + return fallback_idx; +} + +TrtModule& Smollm3TextGenerationPipeline::bind_decoder_for_step() { + const int32_t desired_rows = std::max(state_->preferred_cache_rows(), 1); + const int32_t next_idx = select_decoder_index(desired_rows); + if (!state_bound_ || next_idx != active_decoder_index_) { + active_decoder_index_ = next_idx; + state_->bind_to(*decoders_[static_cast(active_decoder_index_)].module); + state_bound_ = true; + } + return *decoders_[static_cast(active_decoder_index_)].module; +} + +void Smollm3TextGenerationPipeline::run_step(int32_t token_id, std::vector& logits) { + TensorMap inputs; + const int32_t position_before = state_->position(); + const int32_t rows_before = std::max(state_->preferred_cache_rows(), 1); + + Tensor token_tensor; + token_tensor.data = &token_id; + token_tensor.shape = {1}; + token_tensor.dtype = DType::kInt32; + inputs[config_.token_id_name] = token_tensor; + + TrtModule& decoder = bind_decoder_for_step(); + state_->prepare_step(inputs); + + TensorMap outputs = decoder.forward(inputs); + + auto it = outputs.find(logits_output_name_); + if (it == outputs.end()) { + throw std::runtime_error("Smollm3TextGenerationPipeline: no '" + logits_output_name_ + + "' output"); + } + + const auto& logits_tensor = it->second; + auto num_logits = logits_tensor.numel(); + logits.resize(static_cast(num_logits)); + std::memcpy(logits.data(), logits_tensor.data, num_logits * sizeof(float)); + + state_->advance(); + maybe_append_step_trace(position_before, token_id, active_decoder_index_, rows_before, + std::max(state_->preferred_cache_rows(), 1), logits); +} + +void Smollm3TextGenerationPipeline::run_step_device(int32_t token_id) { + TensorMap inputs; + + Tensor token_tensor; + token_tensor.data = &token_id; + token_tensor.shape = {1}; + token_tensor.dtype = DType::kInt32; + inputs[config_.token_id_name] = token_tensor; + + TrtModule& decoder = bind_decoder_for_step(); + state_->prepare_step(inputs); + + // Use forward_async + sync instead of forward() to skip the D2H output copy. + // The GPU argmax kernel reads logits directly from the device buffer. + decoder.forward_async(inputs); + decoder.sync(); + + // Get device pointer to logits output buffer (still on GPU). + d_logits_ptr_ = static_cast(decoder.device_ptr(logits_output_name_)); + + state_->advance(); +} + +int32_t Smollm3TextGenerationPipeline::argmax(const std::vector& logits) { + if (logits.empty()) + return 0; + return static_cast( + std::distance(logits.begin(), std::max_element(logits.begin(), logits.end()))); +} + +} // namespace trtmc diff --git a/src/runtime/models/smollm3/pipeline.h b/src/runtime/models/smollm3/pipeline.h new file mode 100644 index 0000000000..e0a66ef4fa --- /dev/null +++ b/src/runtime/models/smollm3/pipeline.h @@ -0,0 +1,209 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +// Model-owned decoder text pipeline. +// +// Composes: TrtModule (decoder) + Smollm3KvCache + ITokenizer for this runtime +// plugin. Architecture-specific behavior remains in this model directory and +// in the TRT engine emitted by the matching family builder. + +#include "runtime/models/smollm3/inference_state.h" +#include "runtime/models/smollm3/sampler.h" +#include "trtmc/pipeline.h" +#include "trtmc/runtime/trt_module.h" +#include "trtmc/tokenizer.h" + +#include +#include +#include +#include + +namespace trtmc { + +class Smollm3KvCache; + +struct Smollm3TextGenConfig { + int32_t vocab_size{0}; + int32_t id_bos{0}; + int32_t id_eos{0}; + std::vector id_eos_ids; + bool has_position_input{true}; + std::string chat_template_format{}; + std::string token_id_name{"token_id"}; + std::string logits_output_name{"logits"}; + // runtime.* namespace (replaces TRTMC_DISABLE_CUDA_GRAPH, TRTMC_GPU_ARGMAX). + // decoder_plugin::create() populates these from ctx.runtime_config. + bool disable_cuda_graph{false}; + bool prefer_gpu_greedy{false}; + bool log_runtime_stats{false}; + + // Batched-prefill plumbing — populated when the bundle ships with a + // dedicated prefill optimization profile. The runtime forwards the + // prompt through `prefill_module` in one or more profile-bounded chunks + // before falling back to the per-token decode loop. + std::string present_k_pattern{"present_k_{i}"}; + std::string present_v_pattern{"present_v_{i}"}; + int32_t prefill_max_length{0}; + int32_t prefill_profile_index{-1}; + std::string prefill_log_label; + int32_t num_layers{0}; + int32_t kv_dim{0}; + int32_t mask_token_id{-1}; + int32_t diffusion_block_length{32}; + bool supports_text_diffusion{false}; +}; + +// Populate the process-wide step-trace state from the resolved ConfigBundle. +// Called by decoder_plugin::create() before constructing the pipeline. +// Replaces the TRTMC_TEXT_STEP_TRACE_* env vars (deleted). Empty `path` +// keeps tracing disabled; a non-empty path truncates the target file. +void apply_text_trace_config_from_registry(const std::string& path, std::int32_t start_position, + std::int32_t end_position, std::int32_t top_k); + +class Smollm3TextGenerationPipeline final : public IPipeline { + public: + struct DecoderContext { + int32_t kv_rows{0}; + std::unique_ptr module; + }; + + Smollm3TextGenerationPipeline(std::unique_ptr decoder, + std::unique_ptr state, + Smollm3TextGenConfig config, cudaStream_t stream, + std::shared_ptr tokenizer = nullptr, + std::string model_id_str = "", + std::unique_ptr sampler = nullptr, + std::shared_ptr distributed_owner = nullptr); + Smollm3TextGenerationPipeline(std::vector decoders, + std::unique_ptr state, + Smollm3TextGenConfig config, cudaStream_t stream, + std::shared_ptr tokenizer = nullptr, + std::string model_id_str = "", + std::unique_ptr sampler = nullptr, + std::unique_ptr prefill = nullptr, + std::unique_ptr linear_spec_lora_prefill = nullptr, + std::shared_ptr distributed_owner = nullptr); + + // Public API: takes raw text, returns typed result. + TextResult generate(const std::string& prompt, const GenerateConfig& cfg = {}) override; + + const char* model_id() const override { return model_id_.c_str(); } + const char* pipeline_type() const override { return "Smollm3TextGenerationPipeline"; } + + // Token-ID-based generation (for unit tests and internal callers). + struct GenerationResult { + std::vector token_ids; + }; + GenerationResult generate_ids(const std::vector& input_ids, const GenerateConfig& cfg); + + // Argmax over logits (public for testing). + static int32_t argmax(const std::vector& logits); + + private: + // Kept before TRT modules so TP communicators outlive contexts/engines. + std::shared_ptr distributed_owner_; + std::vector decoders_; + std::unique_ptr prefill_; + std::unique_ptr linear_spec_lora_prefill_; + std::unique_ptr state_; + Smollm3TextGenConfig config_; + cudaStream_t stream_; + std::shared_ptr tokenizer_; + std::string model_id_; + std::unique_ptr sampler_; + bool prefer_gpu_greedy_{false}; + const float* d_logits_ptr_{nullptr}; // device logits pointer (for GPU sampling) + std::string logits_output_name_; + int32_t active_decoder_index_{-1}; + bool state_bound_{false}; + double last_setup_ms_{0.0}; + + // Internal: generate from token IDs with sampling parameters and timing. + struct TimedGenResult { + std::vector token_ids; + double prefill_ms{0.0}; + double decode_ms{0.0}; + }; + TimedGenResult generate_from_ids(const std::vector& input_ids, int32_t max_new_tokens, + const Smollm3SamplingParams& params, + const GenerateConfig& cfg); + TimedGenResult generate_diffusion_from_ids(const std::vector& input_ids, + int32_t max_new_tokens, + const Smollm3SamplingParams& params, + const GenerateConfig& cfg); + TimedGenResult generate_linear_spec_from_ids(const std::vector& input_ids, + int32_t max_new_tokens, + const Smollm3SamplingParams& params, + const GenerateConfig& cfg, bool use_lora_draft); + std::string resolve_generation_mode(const GenerateConfig& cfg) const; + void reset_generation_context(); + TrtModule& require_block_prefill(int32_t sq, TrtModule* prefill_override); + Smollm3KvCache& require_block_kv_cache(); + void copy_block_logits(const TensorMap& outputs, std::vector& logits) const; + void append_prefill_kv(Smollm3KvCache& kv, TrtModule& prefill, int32_t sq); + int32_t resolve_text_diffusion_block_length(const GenerateConfig& cfg, int32_t max_new_tokens, + bool require_divisible) const; + int32_t seed_next_token_from_prefill(const std::vector& input_ids, + std::vector& logits, int32_t vocab); + void fill_diffusion_block(std::vector& block, std::vector& logits, + int32_t block_len, int32_t vocab, bool use_threshold, + float threshold); + int32_t verify_diffusion_block(const std::vector& block, std::vector& logits, + int32_t block_len, int32_t vocab); + bool append_tokens_until_eos(const std::vector& tokens, std::vector& output, + const Smollm3SamplingParams& params) const; + void fill_linear_spec_block(std::vector& block, std::vector& logits, + int32_t block_len, int32_t vocab, bool threshold_enabled, + float threshold, bool use_lora_draft); + std::vector verify_linear_spec_block(const std::vector& block, + std::vector& logits, int32_t block_len, + int32_t vocab); + static int32_t count_linear_spec_accepts(const std::vector& ar_tokens, + const std::vector& block); + bool append_linear_spec_tokens(const std::vector& ar_tokens, int32_t emit_count, + std::vector& output, int32_t& generated, + const Smollm3SamplingParams& params) const; + + // Run one decoder step: token_id → logits (D2H to host). Updates cache. + void run_step(int32_t token_id, std::vector& logits); + + // Run one decoder step: logits stay on device (d_logits_ptr_ updated). + void run_step_device(int32_t token_id); + + // Decode loop (extracted for CCN). + int32_t run_decode_loop(Smollm3ISampler* sampler, const Smollm3SamplingParams& params, + std::vector& output, std::vector& logits, + int32_t max_new_tokens, bool gpu_sampling, const GenerateConfig& cfg, + int32_t prompt_token_count); + int32_t select_decoder_index(int32_t desired_rows) const; + TrtModule& bind_decoder_for_step(); + + std::unique_ptr make_step_sampler(const Smollm3SamplingParams& params); + void run_prefill(const std::vector& input_ids, std::vector& logits, + bool gpu_sampling); + void run_prefill_block(const std::vector& input_ids, bool bidirectional, + bool append_kv, std::vector& logits, + TrtModule* prefill_override = nullptr); + // Returns true if the batched prefill engine handled the prompt; false + // means caller must fall back to the per-token decode loop. + bool run_prefill_batched(const std::vector& input_ids, std::vector& logits, + bool retain_device_logits); + void run_prefill_chunk(const int32_t* token_ids, int32_t chunk_size, + const std::vector& present_k, + const std::vector& present_v, Smollm3KvCache& kv, + std::vector& logits, bool retain_device_logits); + void log_batched_prefill(int32_t token_count, int32_t chunk_count, + int32_t max_chunk_size) const; + const TrtModule* generation_capacity_module() const; + void prime_decoder_after_batched_prefill(const std::vector& input_ids); + bool should_stop_on_answer(const std::vector& output, int32_t prompt_token_count, + const GenerateConfig& cfg, int32_t steps, int32_t stop_interval, + bool is_eos) const; + void log_decode_summary(int32_t steps, double ms) const; +}; + +} // namespace trtmc diff --git a/src/runtime/models/smollm3/plugin.cpp b/src/runtime/models/smollm3/plugin.cpp new file mode 100644 index 0000000000..9db3f552c1 --- /dev/null +++ b/src/runtime/models/smollm3/plugin.cpp @@ -0,0 +1,596 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +// DecoderPlugin: handles this model-owned decoder runtime strategy. +// Standard attention-based decoder with device-resident KV cache. + +#include "plugin_helpers.h" +#include "runtime/models/smollm3/chat_templates.h" +#include "runtime/models/smollm3/pipeline.h" +#include "runtime/models/smollm3/tensor_names.h" +#include "runtime/models/smollm3/triattention_kv_cache.h" +#include "trtmc/config/config_bundle.h" +#include "trtmc/runtime/distributed_runtime.h" +#include "trtmc/runtime/pipeline_registry.h" +#include "utils/json_helpers.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace trtmc { + +namespace { + +struct KvCacheRuntimeSizing { + int32_t runtime_rows{0}; + std::uint64_t row_bytes{0}; + std::uint64_t cache_bytes{0}; + bool override_applied{false}; + bool clamped_to_bundle_max{false}; +}; + +struct TensorParallelRuntimeConfig { + bool enabled{false}; + int32_t tp_size{1}; +}; + +struct TensorParallelRuntime { + TensorParallelRuntimeConfig config; + DistributedRuntimeGroup group; +}; + +int32_t dim_at(const std::vector& shape, int32_t dim) { + if (dim < 0 || static_cast(dim) >= shape.size()) + return -1; + const int64_t value = shape[static_cast(dim)]; + if (value <= 0 || value > std::numeric_limits::max()) + return -1; + return static_cast(value); +} + +int32_t cache_row_dim_from_module(const TrtModule& module, const std::string& tensor_name) { + const auto row_dim = [](const std::vector& shape) -> int32_t { + if (shape.size() == 2) + return dim_at(shape, 1); + if (shape.size() == 4) { + const int32_t heads = dim_at(shape, 1); + const int32_t head_dim = dim_at(shape, 3); + if (heads > 0 && head_dim > 0 && + heads <= std::numeric_limits::max() / head_dim) { + return heads * head_dim; + } + } + return -1; + }; + + const int32_t static_dim = row_dim(module.tensor_shape(tensor_name)); + if (static_dim > 0) + return static_dim; + const int32_t profile_count = module.optimization_profile_count(); + for (int32_t profile_idx = 0; profile_idx < profile_count; ++profile_idx) { + const int32_t profile_dim = row_dim( + module.input_profile_shape(tensor_name, profile_idx, ProfileShapeSelector::kMax)); + if (profile_dim > 0) + return profile_dim; + } + throw std::runtime_error("Unable to infer KV row width from engine tensor '" + tensor_name + + "'"); +} + +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); +} + +std::string format_bytes(std::uint64_t bytes) { + std::ostringstream oss; + constexpr double kGiB = 1024.0 * 1024.0 * 1024.0; + constexpr double kMiB = 1024.0 * 1024.0; + oss.setf(std::ios::fixed); + oss.precision(2); + if (bytes >= static_cast(kGiB)) { + oss << (static_cast(bytes) / kGiB) << " GiB"; + return oss.str(); + } + if (bytes >= static_cast(kMiB)) { + oss << (static_cast(bytes) / kMiB) << " MiB"; + return oss.str(); + } + oss.unsetf(std::ios::floatfield); + oss.precision(6); + oss << bytes << " B"; + return oss.str(); +} + +bool engine_uses_native_kv_updates(const TrtModule& module, const Smollm3KvCacheNames& kv_names) { + const bool has_write_indices = module.has_input(kv_names.cache_write_indices); + const bool has_kv_lengths = module.has_input(kv_names.key_value_lengths); + if (has_write_indices != has_kv_lengths) { + throw std::runtime_error( + "SmolLM3 native KV engine must expose both cache_write_indices and " + "key_value_lengths"); + } + return has_write_indices; +} + +void validate_native_kv_marker(const std::string& config_json, bool engine_uses_native_kv) { + const bool declares_native_kv = extract_json_bool(config_json, "native_kv_cache", false); + const bool has_version = + config_json.find("\"native_kv_contract_version\"") != std::string::npos; + if (!declares_native_kv && !has_version && !engine_uses_native_kv) + return; + if (!declares_native_kv || !engine_uses_native_kv || + extract_json_int(config_json, "native_kv_contract_version", 0) != 1) { + throw std::runtime_error("SmolLM3 native KV metadata does not match the engine contract"); + } +} + +std::uint64_t checked_multiply(std::uint64_t lhs, std::uint64_t rhs) { + if (lhs != 0 && rhs > std::numeric_limits::max() / lhs) + throw std::overflow_error("SmolLM3 native KV byte accounting overflow"); + return lhs * rhs; +} + +void admit_native_kv_allocation(const PipelineContext& ctx, bool native_kv, + const KvCacheRuntimeSizing& sizing) { + if (!native_kv) + return; + + std::size_t free_bytes = 0; + std::size_t total_bytes = 0; + const cudaError_t status = cudaMemGetInfo(&free_bytes, &total_bytes); + if (status != cudaSuccess) { + throw std::runtime_error(std::string("SmolLM3 native KV CUDA memory query failed: ") + + cudaGetErrorString(status)); + } + + constexpr std::uint64_t kTwoGiB = 2ULL << 30; + const auto free = static_cast(free_bytes); + const auto total = static_cast(total_bytes); + const auto reserve = std::max(kTwoGiB, total / 10); + const auto available = free > reserve ? free - reserve : 0; + if (sizing.cache_bytes > available) { + throw std::runtime_error( + "SmolLM3 native KV cache admission failed before allocation: capacity=" + + std::to_string(ctx.config.max_cache_length) + + " tokens, required=" + format_bytes(sizing.cache_bytes) + + ", free=" + format_bytes(free) + ", reserve=" + format_bytes(reserve)); + } +} + +void reject_native_kv_size_override(const PipelineContext& ctx) { + if (ctx.kv_cache_size_bytes != 0) { + throw std::invalid_argument( + "SmolLM3 native TensorRT KV cache allocates the model's complete fixed " + "capacity; kv_cache_size_bytes is not supported"); + } +} + +void apply_runtime_kv_size_override(const PipelineContext& ctx, const TrtModule& module, + const Smollm3KvCacheNames& kv_names, + const Smollm3TriAttentionConfig& tri_cfg, + int32_t bundle_max_rows, KvCacheRuntimeSizing& sizing) { + if (!cache_input_supports_runtime_rows(module, kv_names.cache_k.front())) { + throw std::runtime_error( + "This bundle was not built with runtime-resizable KV cache support. " + "Rebuild with trtmc build --dynamic-kv-cache to use --kv-cache-size."); + } + + const std::uint64_t requested_rows = ctx.kv_cache_size_bytes / sizing.row_bytes; + if (requested_rows == 0) { + throw std::runtime_error("--kv-cache-size is smaller than one KV row (" + + format_bytes(sizing.row_bytes) + ")"); + } + + std::uint64_t runtime_rows = requested_rows; + if (runtime_rows > static_cast(bundle_max_rows)) { + runtime_rows = static_cast(bundle_max_rows); + sizing.clamped_to_bundle_max = true; + } + if (runtime_rows > static_cast(std::numeric_limits::max())) { + throw std::runtime_error("Resolved KV cache rows exceed int32 runtime limits"); + } + + sizing.runtime_rows = static_cast(runtime_rows); + sizing.cache_bytes = runtime_rows * sizing.row_bytes; + sizing.override_applied = true; + + if (tri_cfg.enabled && sizing.runtime_rows < tri_cfg.kv_budget) { + const auto minimum_bytes = static_cast(tri_cfg.kv_budget) * sizing.row_bytes; + throw std::runtime_error( + "--kv-cache-size resolves to " + std::to_string(sizing.runtime_rows) + + " rows, but this TriAttention bundle needs at least " + + std::to_string(tri_cfg.kv_budget) + " rows (" + format_bytes(minimum_bytes) + ")"); + } +} + +KvCacheRuntimeSizing resolve_kv_cache_runtime_sizing( + const PipelineContext& ctx, const TrtModule& module, const Smollm3KvCacheNames& kv_names, + DType cache_dtype, const Smollm3TriAttentionConfig& tri_cfg, int32_t kv_dim, bool native_kv) { + KvCacheRuntimeSizing sizing; + const int32_t bundle_max_rows = ctx.config.max_cache_length; + if (ctx.config.num_layers <= 0 || kv_dim <= 0 || bundle_max_rows <= 0) + throw std::runtime_error("SmolLM3 KV geometry must be positive"); + sizing.row_bytes = checked_multiply( + checked_multiply(checked_multiply(static_cast(ctx.config.num_layers), + static_cast(kv_dim)), + static_cast(dtype_size(cache_dtype))), + 2); + sizing.runtime_rows = bundle_max_rows; + sizing.cache_bytes = + checked_multiply(static_cast(bundle_max_rows), sizing.row_bytes); + + if (native_kv) { + reject_native_kv_size_override(ctx); + return sizing; + } + + if (ctx.kv_cache_size_bytes == 0) + return sizing; + + apply_runtime_kv_size_override(ctx, module, kv_names, tri_cfg, bundle_max_rows, sizing); + return sizing; +} + +void validate_native_kv_runtime(const PipelineContext& ctx, const TrtModule& module, + const Smollm3KvCacheNames& kv_names, DType cache_dtype, + const Smollm3TriAttentionConfig& tri_cfg, + const TensorParallelRuntimeConfig& tp_config, bool native_kv) { + validate_native_kv_marker(ctx.config_json, native_kv); + if (!native_kv) + return; + if (cache_dtype != DType::kBFloat16 || tri_cfg.enabled || tp_config.enabled) { + throw std::runtime_error( + "SmolLM3 native KV requires BF16, single-GPU, non-TriAttention runtime"); + } + + const std::vector expected_shape{1, ctx.config.num_kv_heads, + ctx.config.max_cache_length, 128}; + if (module.tensor_shape(kv_names.cache_k.front()) != expected_shape) { + throw std::runtime_error( + "SmolLM3 native KV requires cache shape [1,num_kv_heads,capacity,128]"); + } +} + +} // namespace + +class DecoderPlugin final : public IPipelinePlugin { + public: + std::unique_ptr create(const PipelineContext& ctx) override { + load_ffi_kernels_from_bundle(ctx.bundle); + apply_text_trace_from_registry(ctx.runtime_config); + + auto tokenizer = create_tokenizer_from_bundle(ctx.bundle); + const auto& io = ctx.config.io_map; + Smollm3KvCacheNames kv_names; + build_kv_names(ctx, io, kv_names); + + const DType cache_dtype = cache_dtype_from_precision(ctx.config.precision); + Smollm3TriAttentionConfig tri_cfg = smollm3_parse_triattention_bundle_config( + ctx.config_json, ctx.config.max_cache_length, ctx.runtime_config); + + TensorParallelRuntime tp_runtime; + tp_runtime.config = parse_tensor_parallel_runtime_config(ctx.config_json); + if (tp_runtime.config.enabled) + tp_runtime.group = initialize_tensor_parallel_group(tp_runtime.config.tp_size); + + const std::string engine_section = tp_runtime.config.enabled + ? tp_engine_section_name(tp_runtime.group.rank) + : std::string("engine_plan"); + auto profile_modules = + load_decoder_profile_modules(ctx, engine_section, nullptr, &tp_runtime); + if (profile_modules.modules.empty()) + throw std::runtime_error("No decoder engine profiles were loaded"); + TrtModule& metadata_module = *profile_modules.modules.front().module; + + const bool native_kv = engine_uses_native_kv_updates(metadata_module, kv_names); + validate_native_kv_runtime(ctx, metadata_module, kv_names, cache_dtype, tri_cfg, + tp_runtime.config, native_kv); + const int32_t kv_dim = cache_row_dim_from_module(metadata_module, kv_names.cache_k.front()); + const auto sizing = resolve_kv_cache_runtime_sizing( + ctx, metadata_module, kv_names, cache_dtype, tri_cfg, kv_dim, native_kv); + + const auto decode_profile_roles = detect_decoder_profile_roles( + metadata_module, io.token_id, kv_names.cache_k.front(), ctx.config.max_cache_length); + + std::unique_ptr prefill_module; + auto decoders = build_decoder_contexts(std::move(profile_modules), sizing.runtime_rows, + decode_profile_roles, prefill_module); + cudaStream_t stream = decoders.front().module->stream(); + + int32_t prefill_profile_idx = decode_profile_roles.prefill_profile_idx; + int32_t prefill_max_length = decode_profile_roles.prefill_max_length; + std::string prefill_log_label; + if (!tp_runtime.config.enabled) { + auto split_prefill_module = + load_split_prefill_module(ctx, stream, io, kv_names, prefill_profile_idx, + prefill_max_length, prefill_log_label); + if (split_prefill_module) + prefill_module = std::move(split_prefill_module); + } + + // Split prefill deserialization can consume additional execution-context + // memory. Admit the KV allocation against the free memory that remains + // after every engine/context needed by this pipeline has been loaded. + admit_native_kv_allocation(ctx, native_kv, sizing); + auto state = + build_inference_state(ctx, sizing, tri_cfg, cache_dtype, kv_dim, kv_names, stream); + log_kv_cache_sizing(ctx, sizing, state.get()); + + Smollm3TextGenConfig tgc; + populate_text_gen_config(ctx, tgc, io, decoders.front(), ctx.runtime_config); + apply_chat_template_format(ctx.bundle, tgc); + // Wire batched prefill. Native TensorRT KV engines update the shared + // aliased cache in place; legacy engines copy prefill outputs into it. + tgc.prefill_max_length = prefill_max_length; + tgc.prefill_profile_index = prefill_profile_idx; + tgc.prefill_log_label = std::move(prefill_log_label); + tgc.num_layers = ctx.config.num_layers; + tgc.kv_dim = kv_dim; + tgc.present_k_pattern = io.present_k_pattern; + tgc.present_v_pattern = io.present_v_pattern; + + return std::make_unique( + std::move(decoders), std::move(state), tgc, stream, std::move(tokenizer), + ctx.bundle.info.model_id, nullptr, std::move(prefill_module), nullptr, + tp_runtime.group.owner); + } + + private: + static std::unique_ptr + load_split_prefill_module(const PipelineContext& ctx, cudaStream_t stream, const IoMap& io, + const Smollm3KvCacheNames& kv_names, int32_t& prefill_profile_idx, + int32_t& prefill_max_length, std::string& prefill_log_label) { + if (find_section(ctx.bundle, "prefill_engine_plan") == nullptr) + return nullptr; + + auto split_prefill_modules = + load_decoder_profile_modules(ctx, "prefill_engine_plan", stream, nullptr); + if (split_prefill_modules.modules.empty()) + return nullptr; + + const auto prefill_roles = + detect_decoder_profile_roles(*split_prefill_modules.modules.front().module, io.token_id, + kv_names.cache_k.front(), ctx.config.max_cache_length); + prefill_profile_idx = prefill_roles.prefill_profile_idx; + prefill_max_length = prefill_roles.prefill_max_length; + auto prefill_module = extract_prefill_module(std::move(split_prefill_modules), + prefill_roles, "prefill_engine_plan"); + if (prefill_module) + prefill_log_label = "prefill engine"; + return prefill_module; + } + + static void apply_text_trace_from_registry(const config::ConfigBundle* cfg) { + if (cfg == nullptr) + return; + try { + apply_text_trace_config_from_registry( + cfg->get("text_trace", "step_trace_path"), + cfg->get("text_trace", "step_trace_start_pos"), + cfg->get("text_trace", "step_trace_end_pos"), + cfg->get("text_trace", "step_trace_topk")); + } catch (const std::exception&) { + // Schema not registered or type mismatch — leave disabled. + } + } + + static BackendProfileModules + load_decoder_profile_modules(const PipelineContext& ctx, const std::string& section_name, + cudaStream_t stream, const TensorParallelRuntime* tp_runtime) { + auto* plan = find_section(ctx.bundle, section_name); + if (plan == nullptr || plan->empty()) + throw std::runtime_error(section_name + " section is missing"); + if (ctx.backend == nullptr) + throw std::runtime_error("No backend loaded"); + + auto profile_rows = extract_json_int_array(ctx.config_json, "dynamic_kv_profile_rows", 16); + const int32_t profile_candidates = + profile_rows.empty() ? 2 : static_cast(profile_rows.size() + 1); + std::vector profile_indices; + profile_indices.reserve(static_cast(profile_candidates)); + for (int32_t i = 0; i < profile_candidates; ++i) + profile_indices.push_back(i); + + ModuleCreateOptions opts; + opts.stream = stream; + opts.runtime_cache_path = ctx.runtime_cache_path.c_str(); + opts.cuda_graphs = ctx.cuda_graphs; + if (tp_runtime != nullptr && tp_runtime->config.enabled) { + opts.distributed_communicator = tp_runtime->group.communicator; + opts.distributed_owner = tp_runtime->group.owner; + } + + const auto t0 = std::chrono::steady_clock::now(); + auto modules = + ctx.backend->create_profile_modules(plan->data(), plan->size(), opts, profile_indices); + const auto t1 = std::chrono::steady_clock::now(); + const double load_ms = std::chrono::duration(t1 - t0).count(); + log_trt_load_timing(section_name.c_str(), load_ms, plan->size()); + for (auto& entry : modules.modules) { + entry.module->set_timing_label(entry.profile_idx == 0 ? section_name + ":profile0" + : section_name + ":decode"); + } + return modules; + } + + static void build_kv_names(const PipelineContext& ctx, const IoMap& io, + Smollm3KvCacheNames& kv_names) { + kv_names.position_id = io.position_id; + kv_names.attention_mask = io.attention_mask; + for (int32_t i = 0; i < ctx.config.num_layers; ++i) { + kv_names.cache_k.push_back(smollm3_expand_layer_name(io.cache_k_pattern, i)); + kv_names.cache_v.push_back(smollm3_expand_layer_name(io.cache_v_pattern, i)); + kv_names.present_k.push_back(smollm3_expand_layer_name(io.present_k_pattern, i)); + kv_names.present_v.push_back(smollm3_expand_layer_name(io.present_v_pattern, i)); + } + } + + static std::unique_ptr + extract_prefill_module(BackendProfileModules profile_modules, + const DecoderProfileRoles& profile_roles, const char* section_name) { + if (profile_roles.prefill_profile_idx < 0) + return nullptr; + for (auto& entry : profile_modules.modules) { + if (entry.profile_idx != profile_roles.prefill_profile_idx) + continue; + entry.module->set_timing_label(std::string(section_name) + ":prefill"); + return std::move(entry.module); + } + return nullptr; + } + + static BackendProfileModule* find_profile_module(BackendProfileModules& profile_modules, + int32_t profile_idx) { + auto found = std::find_if( + profile_modules.modules.begin(), profile_modules.modules.end(), + [&](const BackendProfileModule& entry) { return entry.profile_idx == profile_idx; }); + if (found == profile_modules.modules.end()) + return nullptr; + return &*found; + } + + static void extract_engine_plan_prefill_module(BackendProfileModules& profile_modules, + const DecoderProfileRoles& profile_roles, + std::unique_ptr& prefill_module) { + if (profile_roles.prefill_profile_idx < 0) + return; + auto* entry = find_profile_module(profile_modules, profile_roles.prefill_profile_idx); + if (entry == nullptr || !entry->module) + return; + entry->module->set_timing_label("engine_plan:prefill"); + prefill_module = std::move(entry->module); + } + + static std::vector + build_decoder_contexts(BackendProfileModules profile_modules, int32_t runtime_rows, + const DecoderProfileRoles& profile_roles, + std::unique_ptr& prefill_module) { + std::vector decoders; + decoders.reserve(profile_modules.modules.size()); + std::vector available_rows; + available_rows.reserve(profile_roles.decode_profiles.size()); + for (const auto& profile : profile_roles.decode_profiles) + available_rows.push_back(profile.kv_rows); + const auto selected_rows = select_decoder_profile_rows(available_rows, runtime_rows); + for (std::size_t index = 0; index < selected_rows.size(); ++index) { + const auto& profile = profile_roles.decode_profiles[index]; + auto* found = find_profile_module(profile_modules, profile.profile_idx); + if (found == nullptr || !found->module) + continue; + found->module->set_timing_label("engine_plan:decode"); + decoders.push_back(Smollm3TextGenerationPipeline::DecoderContext{ + profile.kv_rows, std::move(found->module)}); + } + + extract_engine_plan_prefill_module(profile_modules, profile_roles, prefill_module); + + if (decoders.empty()) + throw std::runtime_error("No decoder profile available for engine_plan"); + if (decoders.back().kv_rows < runtime_rows) { + throw std::runtime_error( + "Loaded decoder profiles do not cover the selected runtime KV capacity"); + } + return decoders; + } + + static std::unique_ptr + build_inference_state(const PipelineContext& ctx, const KvCacheRuntimeSizing& sizing, + Smollm3TriAttentionConfig& tri_cfg, DType cache_dtype, int32_t kv_dim, + Smollm3KvCacheNames& kv_names, cudaStream_t stream) { + std::unique_ptr state; + if (tri_cfg.enabled) { + auto* stats_sec = find_section(ctx.bundle, tri_cfg.stats_section); + if (stats_sec == nullptr || stats_sec->empty()) + throw std::runtime_error("TriAttention stats section is missing: " + + tri_cfg.stats_section); + std::string stats_json(stats_sec->begin(), stats_sec->end()); + Smollm3TriAttentionStats tri_stats = smollm3_parse_triattention_stats_json( + stats_json, ctx.config.num_heads, ctx.config.num_kv_heads, ctx.config.num_layers); + state = std::make_unique( + ctx.config.num_layers, ctx.config.num_kv_heads, sizing.runtime_rows, kv_dim, stream, + std::move(tri_cfg), std::move(tri_stats), cache_dtype, std::move(kv_names)); + } else { + state = + std::make_unique(ctx.config.num_layers, sizing.runtime_rows, kv_dim, + stream, cache_dtype, std::move(kv_names)); + } + if (!state->ok()) + throw std::runtime_error("Failed to create Smollm3KvCache"); + return state; + } + + static void log_kv_cache_sizing(const PipelineContext& ctx, const KvCacheRuntimeSizing& sizing, + Smollm3InferenceState* state) { + std::cerr << "[trtmc] KV cache rows=" << sizing.runtime_rows + << " (bundle max=" << ctx.config.max_cache_length + << ", row=" << format_bytes(sizing.row_bytes) + << ", cache=" << format_bytes(sizing.cache_bytes) << ", state=" + << format_bytes(static_cast(state->device_memory_bytes())) << ")"; + if (sizing.override_applied) { + std::cerr << " [requested=" << format_bytes(ctx.kv_cache_size_bytes) << "]"; + if (sizing.clamped_to_bundle_max) + std::cerr << " [clamped-to-bundle-max]"; + } + std::cerr << '\n'; + } + + static void + populate_text_gen_config(const PipelineContext& ctx, Smollm3TextGenConfig& tgc, const IoMap& io, + const Smollm3TextGenerationPipeline::DecoderContext& first_dec, + const config::ConfigBundle* runtime_config) { + tgc.vocab_size = ctx.config.vocab_size; + tgc.id_bos = ctx.config.id_bos; + tgc.id_eos = ctx.config.id_eos; + tgc.id_eos_ids = ctx.config.id_eos_ids; + tgc.has_position_input = first_dec.module->has_input(io.position_id); + tgc.token_id_name = io.token_id; + tgc.logits_output_name = io.logits; + if (runtime_config == nullptr) + return; + try { + tgc.disable_cuda_graph = runtime_config->get("runtime", "disable_cuda_graph"); + tgc.prefer_gpu_greedy = runtime_config->get("runtime", "prefer_gpu_greedy"); + tgc.log_runtime_stats = runtime_config->get("platform", "trt_log_stderr"); + } catch (const std::exception&) { + // Schema not registered — stay at defaults. + } + } + + static void apply_chat_template_format(const BundleFile& bundle, Smollm3TextGenConfig& tgc) { + std::string chat_tpl; + auto* tok_cfg_sec = find_section(bundle, "tokenizer_config.json"); + if (tok_cfg_sec != nullptr && !tok_cfg_sec->empty()) { + const std::string tok_cfg_text(tok_cfg_sec->begin(), tok_cfg_sec->end()); + chat_tpl = extract_json_string(tok_cfg_text, "chat_template", ""); + } + if (chat_tpl.empty()) { + auto* tpl_sec = find_section(bundle, "chat_template.jinja"); + if (tpl_sec != nullptr && !tpl_sec->empty()) + chat_tpl.assign(tpl_sec->begin(), tpl_sec->end()); + } + tgc.chat_template_format = smollm3_detect_chat_template_format(chat_tpl); + } +}; + +REGISTER_PIPELINE_PLUGIN_WITH_MANIFEST(register_smollm3_plugin, DecoderPlugin, + "smollm3_decoder_kv_cache"); + +} // namespace trtmc diff --git a/src/runtime/models/smollm3/plugin_helpers.cpp b/src/runtime/models/smollm3/plugin_helpers.cpp new file mode 100644 index 0000000000..12d772037f --- /dev/null +++ b/src/runtime/models/smollm3/plugin_helpers.cpp @@ -0,0 +1,592 @@ +/* + * 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 +#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); + +int32_t cache_rows_from_shape(const std::vector& shape) { + const std::size_t row_index = shape.size() == 4 ? 2 : 0; + if (shape.empty() || row_index >= shape.size()) + return -1; + const int64_t rows = shape[row_index]; + if (rows <= 0 || rows > std::numeric_limits::max()) + return -1; + return static_cast(rows); +} + +} // 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'; +} + +std::vector select_decoder_profile_rows(const std::vector& ordered_profile_rows, + int32_t runtime_rows) { + if (runtime_rows <= 0 || ordered_profile_rows.empty()) { + throw std::invalid_argument( + "Decoder profile selection requires positive runtime rows and profiles"); + } + int32_t previous_rows = 0; + for (const int32_t profile_rows : ordered_profile_rows) { + if (profile_rows <= 0 || profile_rows < previous_rows) { + throw std::invalid_argument("Decoder KV profile rows must be positive and ordered"); + } + previous_rows = profile_rows; + } + + std::vector selected; + selected.reserve(ordered_profile_rows.size()); + for (const int32_t profile_rows : ordered_profile_rows) { + selected.push_back(profile_rows); + if (profile_rows >= runtime_rows) + break; + } + if (selected.back() < runtime_rows) { + throw std::runtime_error("No decoder KV profile can cover the selected runtime capacity"); + } + return selected; +} + +bool cache_input_supports_runtime_rows(const TrtModule& module, const std::string& tensor_name) { + if (!module.input_is_dynamic(tensor_name)) + return false; + const int32_t num_profiles = module.optimization_profile_count(); + for (int32_t profile_idx = 0; profile_idx < num_profiles; ++profile_idx) { + const auto min_shape = + module.input_profile_shape(tensor_name, profile_idx, ProfileShapeSelector::kMin); + const auto max_shape = + module.input_profile_shape(tensor_name, profile_idx, ProfileShapeSelector::kMax); + if (!min_shape.empty() && !max_shape.empty() && min_shape.front() > 0 && + max_shape.front() > min_shape.front()) { + return true; + } + } + return false; +} + +int32_t decoder_profile_cache_rows(const TrtModule& module, const std::string& tensor_name, + int32_t profile_idx, int32_t fallback_rows) { + if (!module.input_is_dynamic(tensor_name)) { + const int32_t static_rows = cache_rows_from_shape(module.tensor_shape(tensor_name)); + if (static_rows > 0) + return static_rows; + } + if (profile_idx >= 0 && profile_idx < module.optimization_profile_count()) { + const int32_t max_rows = cache_rows_from_shape( + module.input_profile_shape(tensor_name, profile_idx, ProfileShapeSelector::kMax)); + if (max_rows > 0) + return max_rows; + } + return fallback_rows; +} + +DecoderProfileRoles detect_decoder_profile_roles(const TrtModule& module, + const std::string& token_id_name, + const std::string& cache_k_name, + int32_t fallback_rows) { + const auto token_max_length = [&](int32_t profile_idx) -> int32_t { + const auto shape = + module.input_profile_shape(token_id_name, profile_idx, ProfileShapeSelector::kMax); + if (shape.empty() || shape.front() <= 0 || + shape.front() > std::numeric_limits::max()) { + return -1; + } + return static_cast(shape.front()); + }; + + DecoderProfileRoles roles; + const int32_t num_profiles = module.optimization_profile_count(); + if (num_profiles <= 0) { + roles.decode_profiles.push_back(DecoderProfileInfo{0, fallback_rows}); + return roles; + } + + for (int32_t profile_idx = 0; profile_idx < num_profiles; ++profile_idx) { + const int32_t token_max = token_max_length(profile_idx); + if (token_max > 1) { + if (token_max > roles.prefill_max_length) { + roles.prefill_profile_idx = profile_idx; + roles.prefill_max_length = token_max; + } + continue; + } + roles.decode_profiles.push_back(DecoderProfileInfo{ + profile_idx, + decoder_profile_cache_rows(module, cache_k_name, profile_idx, fallback_rows)}); + } + + if (roles.decode_profiles.empty()) { + const int32_t fallback_profile = + roles.prefill_profile_idx >= 0 ? roles.prefill_profile_idx : 0; + roles.decode_profiles.push_back(DecoderProfileInfo{ + fallback_profile, + decoder_profile_cache_rows(module, cache_k_name, fallback_profile, fallback_rows)}); + } + return roles; +} + +// 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/smollm3/plugin_helpers.h b/src/runtime/models/smollm3/plugin_helpers.h new file mode 100644 index 0000000000..1faf04c7fb --- /dev/null +++ b/src/runtime/models/smollm3/plugin_helpers.h @@ -0,0 +1,164 @@ +/* + * 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 "runtime/models/smollm3/inference_state.h" +#include "runtime/models/smollm3/kv_cache.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 +}; + +struct DecoderProfileInfo { + int32_t profile_idx{0}; + int32_t kv_rows{0}; +}; + +struct DecoderProfileRoles { + int32_t prefill_profile_idx{-1}; + int32_t prefill_max_length{0}; + std::vector decode_profiles; +}; + +// 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); + +// Return whether the engine input can be rebound to a runtime-selected row +// count. Tensor shape reports may already contain the active/opt shape, so +// dynamicity must come from the module's declared input metadata. +bool cache_input_supports_runtime_rows(const TrtModule& module, const std::string& tensor_name); + +// Read a decode profile's KV row ceiling. Dynamic inputs must use that +// profile's kMAX metadata rather than tensor_shape(), which may report the +// currently active positive shape. +int32_t decoder_profile_cache_rows(const TrtModule& module, const std::string& tensor_name, + int32_t profile_idx, int32_t fallback_rows); + +DecoderProfileRoles detect_decoder_profile_roles(const TrtModule& module, + const std::string& token_id_name, + const std::string& cache_k_name, + int32_t fallback_rows); + +// Keep every decode bucket below the requested runtime capacity plus the +// first ceiling bucket that can execute that capacity. +std::vector select_decoder_profile_rows(const std::vector& ordered_profile_rows, + int32_t runtime_rows); + +// 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/src/runtime/models/smollm3/sampler.cpp b/src/runtime/models/smollm3/sampler.cpp new file mode 100644 index 0000000000..568f26b487 --- /dev/null +++ b/src/runtime/models/smollm3/sampler.cpp @@ -0,0 +1,484 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +// Smollm3ISampler implementations: GreedySampler, TopKSampler, and factory. +// +// GreedySampler wraps the same std::max_element logic as the original +// Smollm3TextGenerationPipeline::argmax() and select_argmax_token(), producing +// bit-identical token sequences. +// +// TopKSampler handles temperature, top-k, top-p, and min-p sampling on host +// logits, with internal xorshift64 RNG state. + +#include "runtime/models/smollm3/sampler.h" + +#include "trtmc/pipeline.h" +#if TRTMC_HAS_CUDA_KERNELS +#include "runtime/models/smollm3/argmax_kernel.h" +#include "runtime/models/smollm3/sparse_multinomial_kernel.h" + +#include +#endif + +#include +#include +#include +#include +#include + +namespace trtmc { + +bool smollm3_is_eos_token(const Smollm3SamplingParams& params, int32_t token_id) { + if (!params.eos_token_ids.empty()) { + return std::find(params.eos_token_ids.begin(), params.eos_token_ids.end(), token_id) != + params.eos_token_ids.end(); + } + return params.eos_token_id >= 0 && token_id == params.eos_token_id; +} + +// Shared argmax helper — returns {token_id, logprob} for the highest logit. +static Smollm3SampleResult argmax_over_logits(const float* logits, int32_t vocab_size, + const Smollm3SamplingParams& params) { + Smollm3SampleResult result; + const float* best = logits; + for (int32_t i = 1; i < vocab_size; ++i) { + if (logits[i] > *best) + best = logits + i; + } + result.token_id = static_cast(best - logits); + result.logprob = *best; + result.is_eos = smollm3_is_eos_token(params, result.token_id); + return result; +} + +struct FilteredDistribution { + std::vector indices; + std::vector probs; + int32_t keep{0}; +}; + +static constexpr float kSamplingEpsilon = 1e-6F; + +static float sanitized_temperature(float temperature) { + if (!std::isfinite(temperature)) + return 1.0F; + return std::max(temperature, 0.0F); +} + +static float sanitized_top_p(float top_p) { + if (!std::isfinite(top_p)) + return 1.0F; + return std::min(std::max(top_p, 0.0F), 1.0F); +} + +static float sanitized_min_p(float min_p) { + if (!std::isfinite(min_p)) + return 0.0F; + return std::min(std::max(min_p, 0.0F), 1.0F); +} + +static bool top_p_enabled(float top_p) { + return top_p > 0.0F && top_p < 1.0F - kSamplingEpsilon; +} + +static bool greedy_equivalent(const Smollm3SamplingParams& params) { + const float temperature = sanitized_temperature(params.temperature); + const float top_p = sanitized_top_p(params.top_p); + return temperature < kSamplingEpsilon || top_p <= 0.0F; +} + +static void topk_indices_and_softmax(FilteredDistribution& dist, const float* logits, int32_t n, + int32_t k, float temperature) { + dist.indices.resize(static_cast(n)); + std::iota(dist.indices.begin(), dist.indices.end(), 0); + std::partial_sort(dist.indices.begin(), dist.indices.begin() + k, dist.indices.end(), + [&](int32_t a, int32_t b) { + return logits[static_cast(a)] > + logits[static_cast(b)]; + }); + const float max_logit = logits[static_cast(dist.indices[0])]; + dist.probs.resize(static_cast(k)); + float sum = 0.0F; + for (int32_t i = 0; i < k; ++i) { + const float scaled = + (logits[static_cast(dist.indices[static_cast(i)])] - + max_logit) / + temperature; + dist.probs[static_cast(i)] = std::isfinite(scaled) ? std::exp(scaled) : 0.0F; + sum += dist.probs[static_cast(i)]; + } + if (sum > 0.0F) { + for (int32_t i = 0; i < k; ++i) + dist.probs[static_cast(i)] /= sum; + } else { + const float uniform = 1.0F / static_cast(k); + for (int32_t i = 0; i < k; ++i) + dist.probs[static_cast(i)] = uniform; + } +} + +static int32_t apply_min_p(const FilteredDistribution& dist, int32_t k, float min_p) { + if (min_p <= 0.0F) + return k; + const float max_prob = dist.probs.empty() ? 1.0F : dist.probs[0]; + if (max_prob <= 0.0F) + return k; + const float min_prob = min_p * max_prob; + int32_t keep = 0; + while (keep < k && dist.probs[static_cast(keep)] >= min_prob) + ++keep; + return std::max(keep, 1); +} + +static int32_t apply_top_p(const FilteredDistribution& dist, int32_t keep, float top_p) { + if (!top_p_enabled(top_p)) + return keep; + float cumulative = 0.0F; + int32_t top_p_keep = 0; + while (top_p_keep < keep) { + cumulative += dist.probs[static_cast(top_p_keep)]; + ++top_p_keep; + if (cumulative >= top_p) + break; + } + return std::max(top_p_keep, 1); +} + +static void renormalize_kept_prefix(FilteredDistribution& dist, int32_t keep) { + float kept_sum = 0.0F; + for (int32_t i = 0; i < keep; ++i) + kept_sum += dist.probs[static_cast(i)]; + if (kept_sum > 0.0F) { + for (int32_t i = 0; i < keep; ++i) + dist.probs[static_cast(i)] /= kept_sum; + } else { + const float uniform = 1.0F / static_cast(keep); + for (int32_t i = 0; i < keep; ++i) + dist.probs[static_cast(i)] = uniform; + } +} + +static FilteredDistribution build_filtered_distribution(const float* logits, int32_t vocab_size, + const Smollm3SamplingParams& params) { + const int32_t n = vocab_size; + const float temperature = sanitized_temperature(params.temperature); + const float top_p = sanitized_top_p(params.top_p); + const float min_p = sanitized_min_p(params.min_p); + const bool full_vocab_for_top_p = top_p_enabled(top_p) && params.top_k <= 1; + const int32_t k = + (params.top_k <= 0 || full_vocab_for_top_p) ? n : std::min(std::max(params.top_k, 1), n); + FilteredDistribution dist; + topk_indices_and_softmax(dist, logits, n, k, temperature); + int32_t keep = apply_min_p(dist, k, min_p); + keep = apply_top_p(dist, keep, top_p); + if (keep < k) + renormalize_kept_prefix(dist, keep); + dist.keep = keep; + return dist; +} + +// ───────────────────────────────────────────────────────────── +// GreedySampler: deterministic argmax (identical to select_argmax_token) +// ───────────────────────────────────────────────────────────── + +class GreedySampler final : public Smollm3ISampler { + public: + Smollm3SampleResult sample(const float* logits, int32_t vocab_size, + const Smollm3SamplingParams& params) override { + if (vocab_size <= 0 || logits == nullptr) { + Smollm3SampleResult result; + result.token_id = 0; + result.is_eos = smollm3_is_eos_token(params, 0); + return result; + } + + return argmax_over_logits(logits, vocab_size, params); + } + + Smollm3LogitsLocation logits_location() const override { return Smollm3LogitsLocation::HOST; } + const char* sampler_type() const override { return "greedy"; } +}; + +// ───────────────────────────────────────────────────────────── +// TopKSampler: temperature-scaled top-k with xorshift64 RNG +// (identical logic to sample_token_topk) +// ───────────────────────────────────────────────────────────── + +class TopKSampler final : public Smollm3ISampler { + public: + explicit TopKSampler(uint64_t initial_seed) + : rng_state_(initial_seed == 0 ? 1 : initial_seed), + initial_seed_(initial_seed == 0 ? 1 : initial_seed) {} + + Smollm3SampleResult sample(const float* logits, int32_t vocab_size, + const Smollm3SamplingParams& params) override { + Smollm3SampleResult result; + + if (vocab_size <= 0 || logits == nullptr) { + result.token_id = 0; + result.is_eos = smollm3_is_eos_token(params, 0); + return result; + } + + if (greedy_equivalent(params)) + return argmax_over_logits(logits, vocab_size, params); + + const FilteredDistribution dist = build_filtered_distribution(logits, vocab_size, params); + + // xorshift64 random number generation + rng_state_ ^= rng_state_ << 13; + rng_state_ ^= rng_state_ >> 7; + rng_state_ ^= rng_state_ << 17; + float u = static_cast(rng_state_ & 0xFFFFFFFF) / 4294967296.0F; + + // Sample from cumulative distribution + float cumulative = 0.0F; + for (int32_t i = 0; i < dist.keep; ++i) { + cumulative += dist.probs[static_cast(i)]; + if (u < cumulative) { + result.token_id = dist.indices[static_cast(i)]; + result.logprob = std::log(std::max(dist.probs[static_cast(i)], + std::numeric_limits::min())); + result.is_eos = smollm3_is_eos_token(params, result.token_id); + return result; + } + } + + result.token_id = dist.indices[static_cast(dist.keep - 1)]; + result.logprob = std::log(std::max(dist.probs[static_cast(dist.keep - 1)], + std::numeric_limits::min())); + result.is_eos = smollm3_is_eos_token(params, result.token_id); + return result; + } + + Smollm3LogitsLocation logits_location() const override { return Smollm3LogitsLocation::HOST; } + const char* sampler_type() const override { return "top_k"; } + + void reset() override { rng_state_ = initial_seed_; } + + private: + uint64_t rng_state_; + uint64_t initial_seed_; +}; + +#if TRTMC_HAS_LIBTORCH_MULTINOMIAL && TRTMC_HAS_CUDA_KERNELS +class TorchCudaMultinomialSampler final : public Smollm3ISampler { + public: + explicit TorchCudaMultinomialSampler(uint64_t initial_seed) + : initial_seed_(initial_seed == 0 ? 1 : initial_seed) { + cudaMalloc(&d_token_id_, sizeof(int32_t)); + } + + ~TorchCudaMultinomialSampler() override { + cudaFree(d_indices_); + cudaFree(d_probs_); + cudaFree(d_token_id_); + } + + TorchCudaMultinomialSampler(const TorchCudaMultinomialSampler&) = delete; + TorchCudaMultinomialSampler& operator=(const TorchCudaMultinomialSampler&) = delete; + + Smollm3SampleResult sample(const float* logits, int32_t vocab_size, + const Smollm3SamplingParams& params) override { + Smollm3SampleResult result; + + if (vocab_size <= 0 || logits == nullptr) { + result.token_id = 0; + result.is_eos = smollm3_is_eos_token(params, 0); + return result; + } + + if (greedy_equivalent(params)) + return argmax_over_logits(logits, vocab_size, params); + + const FilteredDistribution dist = build_filtered_distribution(logits, vocab_size, params); + ensure_execution_policy(vocab_size); + ensure_device_buffers(dist.keep); + + cudaMemcpyAsync(d_indices_, dist.indices.data(), + static_cast(dist.keep) * sizeof(int32_t), + cudaMemcpyHostToDevice, stream_); + cudaMemcpyAsync(d_probs_, dist.probs.data(), + static_cast(dist.keep) * sizeof(float), cudaMemcpyHostToDevice, + stream_); + smollm3_gpu_sparse_torch_multinomial_exact(d_indices_, d_probs_, dist.keep, initial_seed_, + current_offset_, total_threads_, d_token_id_, + stream_); + cudaMemcpyAsync(&h_token_id_, d_token_id_, sizeof(int32_t), cudaMemcpyDeviceToHost, + stream_); + cudaStreamSynchronize(stream_); + current_offset_ += counter_offset_; + + result.token_id = h_token_id_; + float picked_prob = 0.0F; + for (int32_t i = 0; i < dist.keep; ++i) { + if (dist.indices[static_cast(i)] == result.token_id) { + picked_prob = dist.probs[static_cast(i)]; + break; + } + } + result.logprob = std::log(std::max(picked_prob, std::numeric_limits::min())); + result.is_eos = smollm3_is_eos_token(params, result.token_id); + return result; + } + + Smollm3LogitsLocation logits_location() const override { return Smollm3LogitsLocation::HOST; } + const char* sampler_type() const override { return "torch_multinomial"; } + + void reset() override { current_offset_ = 0; } + + private: + void ensure_device_buffers(int32_t keep) { + if (keep > capacity_) { + cudaFree(d_indices_); + cudaFree(d_probs_); + cudaMalloc(&d_indices_, static_cast(keep) * sizeof(int32_t)); + cudaMalloc(&d_probs_, static_cast(keep) * sizeof(float)); + capacity_ = keep; + } + } + + void ensure_execution_policy(int32_t vocab_size) { + if (vocab_size != cached_vocab_size_) { + const Smollm3TorchMultinomialExecutionPolicy policy = + smollm3_compute_torch_multinomial_execution_policy(vocab_size); + cached_vocab_size_ = vocab_size; + total_threads_ = policy.total_threads; + counter_offset_ = policy.counter_offset; + } + } + + uint64_t initial_seed_; + uint64_t current_offset_{0}; + cudaStream_t stream_{nullptr}; + int32_t* d_indices_{nullptr}; + float* d_probs_{nullptr}; + int32_t* d_token_id_{nullptr}; + int32_t h_token_id_{0}; + int32_t capacity_{0}; + int32_t cached_vocab_size_{-1}; + int32_t total_threads_{0}; + uint64_t counter_offset_{0}; +}; +#endif + +// ───────────────────────────────────────────────────────────── +// GpuGreedySampler: on-device argmax (no D2H logit transfer) +// ───────────────────────────────────────────────────────────── + +#if TRTMC_HAS_CUDA_KERNELS +class GpuGreedySampler final : public Smollm3ISampler { + public: + explicit GpuGreedySampler(cudaStream_t stream) : stream_(stream) { + cudaMalloc(&d_token_id_, sizeof(int32_t)); + cudaMalloc(&d_logit_val_, sizeof(float)); + } + + ~GpuGreedySampler() override { + cudaFree(d_token_id_); + cudaFree(d_logit_val_); + } + + GpuGreedySampler(const GpuGreedySampler&) = delete; + GpuGreedySampler& operator=(const GpuGreedySampler&) = delete; + + Smollm3SampleResult sample(const float* logits, int32_t vocab_size, + const Smollm3SamplingParams& params) override { + Smollm3SampleResult result; + if (vocab_size <= 0 || logits == nullptr) { + result.token_id = 0; + result.is_eos = smollm3_is_eos_token(params, 0); + return result; + } + + // logits is a device pointer — run GPU argmax kernel + smollm3_gpu_argmax(logits, vocab_size, d_token_id_, d_logit_val_, stream_); + + // D2H: copy only token_id + logit (8 bytes total vs vocab_size*4 bytes) + cudaMemcpyAsync(&h_token_id_, d_token_id_, sizeof(int32_t), cudaMemcpyDeviceToHost, + stream_); + cudaMemcpyAsync(&h_logit_val_, d_logit_val_, sizeof(float), cudaMemcpyDeviceToHost, + stream_); + cudaStreamSynchronize(stream_); + + result.token_id = h_token_id_; + result.logprob = h_logit_val_; + result.is_eos = smollm3_is_eos_token(params, result.token_id); + return result; + } + + Smollm3LogitsLocation logits_location() const override { return Smollm3LogitsLocation::DEVICE; } + const char* sampler_type() const override { return "gpu_greedy"; } + + private: + cudaStream_t stream_{nullptr}; + int32_t* d_token_id_{nullptr}; + float* d_logit_val_{nullptr}; + int32_t h_token_id_{0}; + float h_logit_val_{0.0f}; +}; +#endif // TRTMC_HAS_CUDA_KERNELS + +// ───────────────────────────────────────────────────────────── +// Factory +// ───────────────────────────────────────────────────────────── + +Smollm3SamplingParams +smollm3_sampling_params_from_config(const GenerateConfig& cfg, + const std::vector& default_eos_token_ids) { + Smollm3SamplingParams p; + p.temperature = cfg.temperature; + p.top_k = cfg.top_k; + p.top_p = cfg.top_p; + p.min_p = cfg.min_p; + p.seed = cfg.seed; + p.eos_token_ids = + (cfg.eos_token_id >= 0) ? std::vector{cfg.eos_token_id} : default_eos_token_ids; + p.eos_token_id = p.eos_token_ids.empty() ? -1 : p.eos_token_ids.front(); + return p; +} + +Smollm3SamplingParams smollm3_sampling_params_from_config(const GenerateConfig& cfg, + int32_t default_eos) { + const std::vector defaults = + default_eos >= 0 ? std::vector{default_eos} : std::vector{}; + return smollm3_sampling_params_from_config(cfg, defaults); +} + +std::unique_ptr +create_smollm3_sampler(const Smollm3SamplingParams& params, + [[maybe_unused]] const Smollm3SamplerFactoryOptions& options) { + // Greedy when sampling is fully disabled and no explicit random seed is set. + const float top_p = sanitized_top_p(params.top_p); + const float min_p = sanitized_min_p(params.min_p); + if (params.top_k <= 1 && top_p >= 1.0F - kSamplingEpsilon && min_p <= 0.0F && params.seed < 0) { + return std::make_unique(); + } + + uint64_t seed = (params.seed >= 0) ? static_cast(params.seed) + : 42ULL; // deterministic default for reproducibility +#if TRTMC_HAS_LIBTORCH_MULTINOMIAL && TRTMC_HAS_CUDA_KERNELS + if (options.prefer_torch_cuda_multinomial) + return std::make_unique(seed); +#endif + + // TopK sampler with xorshift64 RNG + return std::make_unique(seed); +} + +std::unique_ptr create_smollm3_sampler(const Smollm3SamplingParams& params) { + return create_smollm3_sampler(params, Smollm3SamplerFactoryOptions{}); +} + +std::unique_ptr create_smollm3_gpu_greedy_sampler(void* stream) { +#if TRTMC_HAS_CUDA_KERNELS + return std::make_unique(static_cast(stream)); +#else + (void)stream; + return nullptr; +#endif +} + +} // namespace trtmc diff --git a/src/runtime/models/smollm3/sampler.h b/src/runtime/models/smollm3/sampler.h new file mode 100644 index 0000000000..daf5bff71b --- /dev/null +++ b/src/runtime/models/smollm3/sampler.h @@ -0,0 +1,107 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +// Smollm3ISampler: token selection abstraction for autoregressive generation. +// +// Decouples token selection strategy (greedy, top-k, nucleus, beam search, +// grammar-constrained, on-device argmax) from the generation loop. +// Pipelines call sampler->sample() without knowing the selection strategy. +// +// Smollm3LogitsLocation tells the pipeline whether to D2H transfer logits before +// calling sample(). HOST samplers (current default) require logits on CPU. +// DEVICE samplers (future: TASK-08) read logits directly from GPU memory. + +#include +#include +#include + +namespace trtmc { + +/// Sampling parameters -- controls token selection behavior. +struct Smollm3SamplingParams { + float temperature{1.0f}; + int32_t top_k{1}; // 1 = greedy unless top_p is active; <=0 = no top-k limit + float top_p{1.0f}; // 1.0 = disabled; 0.0 = greedy; (0,1) = nucleus + float min_p{0.0f}; // 0.0 = disabled; filters tokens below min_p * max_prob + float repetition_penalty{1.0f}; + int32_t seed{-1}; // -1 = deterministic (argmax) + int32_t eos_token_id{-1}; + std::vector eos_token_ids; +}; + +/// Factory options for choosing concrete sampler implementations. +struct Smollm3SamplerFactoryOptions { + bool prefer_torch_cuda_multinomial{true}; +}; + +/// Where the sampler expects logits to live. +enum class Smollm3LogitsLocation { + HOST, // Sampler reads from CPU memory (current default) + DEVICE, // Sampler reads from GPU memory (for on-device sampling) +}; + +/// Token selection result. +struct Smollm3SampleResult { + int32_t token_id{0}; + float logprob{0.0f}; // log-probability of selected token (informational) + bool is_eos{false}; // true if token_id matches eos_token_id +}; + +/// smollm3-owned sampler interface. +class Smollm3ISampler { + public: + virtual ~Smollm3ISampler() = default; + + /// Select the next token from logits. + /// logits: float[vocab_size] on host or device (see logits_location()). + /// vocab_size: number of logit values. + /// params: sampling parameters for this step. + virtual Smollm3SampleResult sample(const float* logits, int32_t vocab_size, + const Smollm3SamplingParams& params) = 0; + + /// Where does this sampler expect logits? + /// HOST: pipeline must D2H logits before calling sample(). + /// DEVICE: pipeline passes device pointer directly (no D2H). + virtual Smollm3LogitsLocation logits_location() const = 0; + + /// Human-readable name for diagnostics. + virtual const char* sampler_type() const = 0; + + /// Reset sampler state (e.g., RNG state between sequences). + virtual void reset() {} +}; + +/// Build Smollm3SamplingParams from GenerateConfig fields. +/// Forward-declared here; defined in sampler.cpp alongside the factory. +struct GenerateConfig; // defined in trtmc/pipeline.h + +Smollm3SamplingParams +smollm3_sampling_params_from_config(const GenerateConfig& cfg, + const std::vector& default_eos_token_ids); +Smollm3SamplingParams smollm3_sampling_params_from_config(const GenerateConfig& cfg, + int32_t default_eos = -1); + +/// Return true when token_id matches any effective EOS token. +bool smollm3_is_eos_token(const Smollm3SamplingParams& params, int32_t token_id); + +/// Factory: create sampler from Smollm3SamplingParams. +/// - top_k <= 1 && top_p/min_p disabled && seed == -1 => GreedySampler +/// - otherwise => TorchCudaMultinomialSampler when compiled in and preferred, +/// falling back to TopKSampler +std::unique_ptr create_smollm3_sampler(const Smollm3SamplingParams& params); +std::unique_ptr +create_smollm3_sampler(const Smollm3SamplingParams& params, + const Smollm3SamplerFactoryOptions& options); + +/// Factory: create a GPU-side greedy sampler (on-device argmax). +/// Requires CUDA kernels (TRTMC_HAS_CUDA_KERNELS). Returns nullptr if unavailable. +/// The sampler reads logits directly from GPU memory and copies back only the +/// token ID (4 bytes) instead of the full logit vector (~600KB for 151K vocab). +/// Pass the CUDA stream used by the pipeline for synchronized kernel execution. +std::unique_ptr create_smollm3_gpu_greedy_sampler(void* stream); + +} // namespace trtmc diff --git a/src/runtime/models/smollm3/sparse_multinomial_kernel.cu b/src/runtime/models/smollm3/sparse_multinomial_kernel.cu new file mode 100644 index 0000000000..643d5ca696 --- /dev/null +++ b/src/runtime/models/smollm3/sparse_multinomial_kernel.cu @@ -0,0 +1,119 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/models/smollm3/sparse_multinomial_kernel.h" + +#include + +#include +#include +#include + +namespace trtmc { + +namespace { + +constexpr int kDistributionBlockSize = 256; +constexpr int kSamplerBlockSize = 128; +constexpr uint64_t kGeneratorOffsetsPerCurandCall = 4; + +__device__ float torch_exponential_from_uniform(float val) { + const float log_val = val >= 1.0F - FLT_EPSILON / 2.0F ? -FLT_EPSILON / 2.0F : logf(val); + return -log_val; +} + +__global__ void sparse_multinomial_exact_kernel(const int32_t* __restrict__ indices, + const float* __restrict__ probs, int32_t keep, + uint64_t seed, uint64_t base_offset, + int32_t total_threads, + int32_t* __restrict__ out_token_id) { + __shared__ float s_scores[kSamplerBlockSize]; + __shared__ int32_t s_tokens[kSamplerBlockSize]; + + const int tid = threadIdx.x; + float best_score = -FLT_MAX; + int32_t best_token = 0; + + for (int32_t i = tid; i < keep; i += blockDim.x) { + const int32_t token_id = indices[i]; + const int64_t linear_index = static_cast(token_id); + const int64_t q = linear_index / total_threads; + const uint64_t loop_iteration = static_cast(q / 4); + const int component = static_cast(q % 4); + const int64_t subsequence = linear_index % total_threads; + + curandStatePhilox4_32_10_t state; + curand_init(seed, static_cast(subsequence), + base_offset + kGeneratorOffsetsPerCurandCall * loop_iteration, &state); + const float4 rand = curand_uniform4(&state); + const float uniform = component == 0 ? rand.x : component == 1 ? rand.y : component == 2 ? rand.z : rand.w; + const float score = probs[i] / torch_exponential_from_uniform(uniform); + if (score > best_score) { + best_score = score; + best_token = token_id; + } + } + + s_scores[tid] = best_score; + s_tokens[tid] = best_token; + __syncthreads(); + + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (tid < stride && s_scores[tid + stride] > s_scores[tid]) { + s_scores[tid] = s_scores[tid + stride]; + s_tokens[tid] = s_tokens[tid + stride]; + } + __syncthreads(); + } + + if (tid == 0) { + *out_token_id = s_tokens[0]; + } +} + +} // namespace + +Smollm3TorchMultinomialExecutionPolicy smollm3_compute_torch_multinomial_execution_policy(int32_t numel) { + if (numel <= 0) { + return {}; + } + + int device = 0; + cudaGetDevice(&device); + cudaDeviceProp props{}; + cudaGetDeviceProperties(&props, device); + + const uint32_t blocks_per_sm = + static_cast(props.maxThreadsPerMultiProcessor / kDistributionBlockSize); + const uint32_t grid = std::min( + static_cast(props.multiProcessorCount) * blocks_per_sm, + static_cast((static_cast(numel) + kDistributionBlockSize - 1) + / kDistributionBlockSize)); + const uint64_t total_threads = static_cast(grid) * kDistributionBlockSize; + const uint64_t counter_offset = + ((static_cast(numel) - 1) + / (total_threads * kGeneratorOffsetsPerCurandCall) + 1) + * kGeneratorOffsetsPerCurandCall; + + Smollm3TorchMultinomialExecutionPolicy policy; + policy.total_threads = static_cast(total_threads); + policy.counter_offset = counter_offset; + return policy; +} + +void smollm3_gpu_sparse_torch_multinomial_exact(const int32_t* d_indices, const float* d_probs, + int32_t keep, uint64_t seed, uint64_t base_offset, + int32_t total_threads, int32_t* d_token_id, + cudaStream_t stream) { + if (keep <= 0 || d_indices == nullptr || d_probs == nullptr || d_token_id == nullptr + || total_threads <= 0) { + return; + } + + sparse_multinomial_exact_kernel<<<1, kSamplerBlockSize, 0, stream>>>( + d_indices, d_probs, keep, seed, base_offset, total_threads, d_token_id); +} + +} // namespace trtmc diff --git a/src/runtime/models/smollm3/sparse_multinomial_kernel.h b/src/runtime/models/smollm3/sparse_multinomial_kernel.h new file mode 100644 index 0000000000..9e01e2a0e8 --- /dev/null +++ b/src/runtime/models/smollm3/sparse_multinomial_kernel.h @@ -0,0 +1,26 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include + +namespace trtmc { + +struct Smollm3TorchMultinomialExecutionPolicy { + int32_t total_threads{0}; + uint64_t counter_offset{0}; +}; + +Smollm3TorchMultinomialExecutionPolicy +smollm3_compute_torch_multinomial_execution_policy(int32_t numel); + +void smollm3_gpu_sparse_torch_multinomial_exact(const int32_t* d_indices, const float* d_probs, + int32_t keep, uint64_t seed, uint64_t base_offset, + int32_t total_threads, int32_t* d_token_id, + cudaStream_t stream); + +} // namespace trtmc diff --git a/src/runtime/models/smollm3/tensor_names.h b/src/runtime/models/smollm3/tensor_names.h new file mode 100644 index 0000000000..18cb7f6e80 --- /dev/null +++ b/src/runtime/models/smollm3/tensor_names.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 +#include + +namespace trtmc { + +inline std::string smollm3_expand_layer_name(const std::string& pattern, int32_t layer) { + std::string result = pattern; + auto replace_all = [&](const std::string& token, const std::string& value) { + std::size_t pos = 0; + while ((pos = result.find(token, pos)) != std::string::npos) { + result.replace(pos, token.size(), value); + pos += value.size(); + } + }; + + replace_all("{2i+2}", std::to_string(2 * layer + 2)); + replace_all("{2i+1}", std::to_string(2 * layer + 1)); + replace_all("{2i}", std::to_string(2 * layer)); + replace_all("{i}", std::to_string(layer)); + return result; +} + +inline std::string smollm3_layer_tensor_name(const char* stem, int32_t layer) { + return std::string(stem) + "_" + std::to_string(layer); +} + +} // namespace trtmc diff --git a/src/runtime/models/smollm3/triattention_kernels.cu b/src/runtime/models/smollm3/triattention_kernels.cu new file mode 100644 index 0000000000..704fc1c5ac --- /dev/null +++ b/src/runtime/models/smollm3/triattention_kernels.cu @@ -0,0 +1,305 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/models/smollm3/triattention_kernels.h" + + +#include +#include +#include + +#include +#include + +namespace trtmc { + +namespace { + +constexpr int kScoreBlockSize = 64; +constexpr int kCompactBlockSize = 256; +constexpr int kWarpSize = 32; +constexpr int kMaxOffsets = 32; +constexpr float kAbsFloor = 1.0e-8F; + +template +__device__ inline float load_as_float(const T* ptr, int32_t idx); + +template <> +__device__ inline float load_as_float(const float* ptr, int32_t idx) { + return ptr[idx]; +} + +template <> +__device__ inline float load_as_float<__half>(const __half* ptr, int32_t idx) { + return __half2float(ptr[idx]); +} + +template <> +__device__ inline float load_as_float<__nv_bfloat16>(const __nv_bfloat16* ptr, int32_t idx) { + return __bfloat162float(ptr[idx]); +} + +template +__device__ inline float load_cache_value(const void* base, int32_t idx) { + const T* typed = static_cast(base); + return load_as_float(typed, idx); +} + +__device__ inline float warp_sum(float value) { + for (int offset = kWarpSize / 2; offset > 0; offset >>= 1) + value += __shfl_down_sync(0xFFFFFFFFU, value, offset); + return value; +} + +template +__global__ void triattention_score_candidates_kernel( + const void* __restrict__ d_cache, int32_t kv_dim, int32_t head_dim, + bool rope_interleaved, const int32_t* __restrict__ candidate_indices, + int32_t candidate_count, const int32_t* __restrict__ positions_per_head, + const float* __restrict__ inv_freq, const float* __restrict__ cos_phase, + const float* __restrict__ sin_phase, int32_t num_offsets, + const int32_t* __restrict__ head_offsets, const int32_t* __restrict__ head_cache_indices, + const float* __restrict__ q_mean_real, const float* __restrict__ q_mean_imag, + const float* __restrict__ q_abs_mean, const float* __restrict__ freq_scale_sq, + int32_t kv_head_count, bool disable_mlr, bool disable_trig, bool aggregation_max, + float* __restrict__ scores_out) { + const int32_t sampled_idx = blockIdx.x; + const int32_t candidate_idx = blockIdx.y; + const int32_t tid = threadIdx.x; + const int32_t warp_id = tid / kWarpSize; + const int32_t lane = tid % kWarpSize; + const int32_t warp_count = (blockDim.x + kWarpSize - 1) / kWarpSize; + + const int32_t half_dim = head_dim / 2; + if (sampled_idx >= kv_head_count || candidate_idx >= candidate_count || num_offsets > kMaxOffsets || + half_dim <= 0) + return; + + __shared__ float shared[(kMaxOffsets + 1) * 4]; + + const int32_t row = candidate_indices[candidate_idx]; + const int32_t head_offset = head_offsets[sampled_idx]; + const int32_t row_base = row * kv_dim + head_offset; + + const float* q_real_row = q_mean_real + sampled_idx * half_dim; + const float* q_imag_row = q_mean_imag + sampled_idx * half_dim; + const float* q_abs_row = q_abs_mean + sampled_idx * half_dim; + const float* freq_scale_row = freq_scale_sq + sampled_idx * half_dim; + + float local_additive = 0.0F; + float local_trig[kMaxOffsets]; + #pragma unroll + for (int o = 0; o < kMaxOffsets; ++o) + local_trig[o] = 0.0F; + + for (int32_t d = tid; d < half_dim; d += blockDim.x) { + float k_rot_real = 0.0F; + float k_rot_imag = 0.0F; + if (rope_interleaved) { + k_rot_real = load_cache_value(d_cache, row_base + 2 * d); + k_rot_imag = load_cache_value(d_cache, row_base + 2 * d + 1); + } else { + k_rot_real = load_cache_value(d_cache, row_base + d); + k_rot_imag = load_cache_value(d_cache, row_base + half_dim + d); + } + + const float q_real = q_real_row[d]; + const float q_imag = q_imag_row[d]; + const float q_abs = q_abs_row[d]; + const float freq_scale_sq_val = freq_scale_row[d]; + const float q_mean_abs = sqrtf(fmaxf(q_real * q_real + q_imag * q_imag, kAbsFloor)); + const float k_abs = sqrtf(fmaxf(k_rot_real * k_rot_real + k_rot_imag * k_rot_imag, kAbsFloor)); + const float prod_real = q_real * k_rot_real + q_imag * k_rot_imag; + const float prod_imag = q_imag * k_rot_real - q_real * k_rot_imag; + const float extra_coef = disable_mlr ? q_abs : (q_abs - q_mean_abs); + local_additive += k_abs * extra_coef * freq_scale_sq_val; + + if (!disable_trig) { + for (int32_t o = 0; o < num_offsets; ++o) { + const int32_t phase_idx = o * half_dim + d; + local_trig[o] += + freq_scale_sq_val + * (prod_real * cos_phase[phase_idx] - prod_imag * sin_phase[phase_idx]); + } + } + } + + local_additive = warp_sum(local_additive); + for (int32_t o = 0; o < num_offsets; ++o) + local_trig[o] = warp_sum(local_trig[o]); + + if (lane == 0) { + const int32_t base = warp_id * (kMaxOffsets + 1); + shared[base] = local_additive; + for (int32_t o = 0; o < num_offsets; ++o) + shared[base + 1 + o] = local_trig[o]; + } + __syncthreads(); + + if (warp_id == 0) { + float block_additive = (lane < warp_count) ? shared[lane * (kMaxOffsets + 1)] : 0.0F; + float block_trig[kMaxOffsets]; + #pragma unroll + for (int o = 0; o < kMaxOffsets; ++o) + block_trig[o] = 0.0F; + for (int32_t o = 0; o < num_offsets; ++o) { + block_trig[o] = + (lane < warp_count) ? shared[lane * (kMaxOffsets + 1) + 1 + o] : 0.0F; + } + + block_additive = warp_sum(block_additive); + for (int32_t o = 0; o < num_offsets; ++o) + block_trig[o] = warp_sum(block_trig[o]); + + if (lane == 0) { + float trig_term = 0.0F; + if (!disable_trig && num_offsets > 0) { + if (aggregation_max) { + trig_term = block_trig[0]; + for (int32_t o = 1; o < num_offsets; ++o) + trig_term = fmaxf(trig_term, block_trig[o]); + } else { + for (int32_t o = 0; o < num_offsets; ++o) + trig_term += block_trig[o]; + trig_term /= static_cast(num_offsets); + } + } + scores_out[sampled_idx * candidate_count + candidate_idx] = trig_term + block_additive; + } + } +} + +template +__global__ void triattention_compact_rows_kernel(const T* __restrict__ src, + T* __restrict__ scratch, int32_t kv_dim, + const int32_t* __restrict__ keep_indices, + int32_t keep_count, int32_t head_group_width, + int32_t num_kv_heads) { + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = + static_cast(keep_count) * static_cast(num_kv_heads) * head_group_width; + if (idx >= total) + return; + + const int32_t dst_row = static_cast(idx / (num_kv_heads * head_group_width)); + const int32_t rem = static_cast(idx % (num_kv_heads * head_group_width)); + const int32_t kv_head = rem / head_group_width; + const int32_t elem_in_group = rem % head_group_width; + const int32_t src_row = keep_indices[kv_head * keep_count + dst_row]; + const int32_t col = kv_head * head_group_width + elem_in_group; + scratch[static_cast(dst_row) * kv_dim + col] = + src[static_cast(src_row) * kv_dim + col]; +} + +template +bool launch_score_kernel(const void* d_cache, int32_t kv_dim, int32_t head_dim, + bool rope_interleaved, + const int32_t* d_candidate_indices, int32_t candidate_count, + const int32_t* d_positions_per_head, const float* d_inv_freq, + const float* d_cos_phase, const float* d_sin_phase, int32_t num_offsets, + const int32_t* d_head_offsets, const int32_t* d_head_cache_indices, + const float* d_q_mean_real, const float* d_q_mean_imag, + const float* d_q_abs_mean, const float* d_freq_scale_sq, + int32_t kv_head_count, + bool disable_mlr, bool disable_trig, bool aggregation_max, + float* d_scores_out, cudaStream_t stream) { + dim3 block(kScoreBlockSize); + dim3 grid(static_cast(kv_head_count), static_cast(candidate_count)); + triattention_score_candidates_kernel<<>>( + d_cache, kv_dim, head_dim, rope_interleaved, d_candidate_indices, candidate_count, + d_positions_per_head, d_inv_freq, d_cos_phase, d_sin_phase, num_offsets, d_head_offsets, + d_head_cache_indices, d_q_mean_real, d_q_mean_imag, d_q_abs_mean, d_freq_scale_sq, + kv_head_count, disable_mlr, disable_trig, aggregation_max, d_scores_out); + return cudaGetLastError() == cudaSuccess; +} + +template +bool launch_compact_kernel(const void* d_src, void* d_scratch, int32_t kv_dim, + const int32_t* d_keep_indices, int32_t keep_count, int32_t head_group_width, + int32_t num_kv_heads, + cudaStream_t stream) { + const int64_t total = + static_cast(keep_count) * static_cast(num_kv_heads) * head_group_width; + const int blocks = static_cast((total + kCompactBlockSize - 1) / kCompactBlockSize); + triattention_compact_rows_kernel<<>>( + static_cast(d_src), static_cast(d_scratch), kv_dim, d_keep_indices, keep_count, + head_group_width, num_kv_heads); + return cudaGetLastError() == cudaSuccess; +} + +} // namespace + +bool smollm3_triattention_score_candidates_gpu( + const void* d_cache, DType cache_dtype, int32_t kv_dim, int32_t head_dim, + bool rope_interleaved, const int32_t* d_candidate_indices, int32_t candidate_count, + const int32_t* d_positions_per_head, const float* d_inv_freq, const float* d_cos_phase, + const float* d_sin_phase, int32_t num_offsets, const int32_t* d_head_offsets, + const int32_t* d_head_cache_indices, const float* d_q_mean_real, + const float* d_q_mean_imag, + const float* d_q_abs_mean, const float* d_freq_scale_sq, int32_t kv_head_count, bool disable_mlr, + bool disable_trig, + bool aggregation_max, float* d_scores_out, cudaStream_t stream) { + if (candidate_count <= 0 || kv_head_count <= 0 || head_dim <= 0) + return false; + if (num_offsets <= 0 || num_offsets > kMaxOffsets) + return false; + (void)d_positions_per_head; + (void)d_inv_freq; + + switch (cache_dtype) { + case DType::kFloat32: + return launch_score_kernel( + d_cache, kv_dim, head_dim, rope_interleaved, d_candidate_indices, candidate_count, + d_positions_per_head, d_inv_freq, d_cos_phase, d_sin_phase, num_offsets, d_head_offsets, + d_head_cache_indices, d_q_mean_real, d_q_mean_imag, + d_q_abs_mean, d_freq_scale_sq, kv_head_count, disable_mlr, disable_trig, aggregation_max, + d_scores_out, stream); + case DType::kFloat16: + return launch_score_kernel<__half>( + d_cache, kv_dim, head_dim, rope_interleaved, d_candidate_indices, candidate_count, + d_positions_per_head, d_inv_freq, d_cos_phase, d_sin_phase, num_offsets, d_head_offsets, + d_head_cache_indices, d_q_mean_real, d_q_mean_imag, + d_q_abs_mean, d_freq_scale_sq, kv_head_count, disable_mlr, disable_trig, aggregation_max, + d_scores_out, stream); + case DType::kBFloat16: + return launch_score_kernel<__nv_bfloat16>( + d_cache, kv_dim, head_dim, rope_interleaved, d_candidate_indices, candidate_count, + d_positions_per_head, d_inv_freq, d_cos_phase, d_sin_phase, num_offsets, d_head_offsets, + d_head_cache_indices, d_q_mean_real, d_q_mean_imag, + d_q_abs_mean, d_freq_scale_sq, kv_head_count, disable_mlr, disable_trig, aggregation_max, + d_scores_out, stream); + default: + return false; + } +} + +bool smollm3_triattention_compact_rows_gpu(const void* d_src, void* d_scratch, DType cache_dtype, + int32_t kv_dim, const int32_t* d_keep_indices, int32_t keep_count, + int32_t head_dim, int32_t num_kv_heads, int32_t query_group_size, + cudaStream_t stream) { + if (keep_count <= 0 || kv_dim <= 0) + return false; + if (head_dim <= 0 || num_kv_heads <= 0 || query_group_size <= 0) + return false; + const int32_t head_group_width = head_dim * query_group_size; + + switch (cache_dtype) { + case DType::kFloat32: + return launch_compact_kernel(d_src, d_scratch, kv_dim, d_keep_indices, keep_count, + head_group_width, num_kv_heads, stream); + case DType::kFloat16: + return launch_compact_kernel<__half>(d_src, d_scratch, kv_dim, d_keep_indices, keep_count, + head_group_width, num_kv_heads, stream); + case DType::kBFloat16: + return launch_compact_kernel<__nv_bfloat16>(d_src, d_scratch, kv_dim, d_keep_indices, + keep_count, head_group_width, num_kv_heads, + stream); + default: + return false; + } +} + +} // namespace trtmc + diff --git a/src/runtime/models/smollm3/triattention_kernels.h b/src/runtime/models/smollm3/triattention_kernels.h new file mode 100644 index 0000000000..8d149640fa --- /dev/null +++ b/src/runtime/models/smollm3/triattention_kernels.h @@ -0,0 +1,31 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "trtmc/runtime/tensor.h" + +#include +#include + +namespace trtmc { + +bool smollm3_triattention_score_candidates_gpu( + const void* d_cache, DType cache_dtype, int32_t kv_dim, int32_t head_dim, bool rope_interleaved, + const int32_t* d_candidate_indices, int32_t candidate_count, + const int32_t* d_positions_per_head, const float* d_inv_freq, const float* d_cos_phase, + const float* d_sin_phase, int32_t num_offsets, const int32_t* d_head_offsets, + const int32_t* d_head_cache_indices, const float* d_q_mean_real, const float* d_q_mean_imag, + const float* d_q_abs_mean, const float* d_freq_scale_sq, int32_t kv_head_count, + bool disable_mlr, bool disable_trig, bool aggregation_max, float* d_scores_out, + cudaStream_t stream); + +bool smollm3_triattention_compact_rows_gpu(const void* d_src, void* d_scratch, DType cache_dtype, + int32_t kv_dim, const int32_t* d_keep_indices, + int32_t keep_count, int32_t head_dim, + int32_t num_kv_heads, int32_t query_group_size, + cudaStream_t stream); + +} // namespace trtmc diff --git a/src/runtime/models/smollm3/triattention_kv_cache.cpp b/src/runtime/models/smollm3/triattention_kv_cache.cpp new file mode 100644 index 0000000000..c2588934cf --- /dev/null +++ b/src/runtime/models/smollm3/triattention_kv_cache.cpp @@ -0,0 +1,2211 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/models/smollm3/triattention_kv_cache.h" + +#include "trtmc/config/config_bundle.h" +#include "trtmc/config/schema_registry.h" +#include "trtmc/runtime/trt_module.h" +#ifdef TRTMC_HAS_CUDA_KERNELS +#include "runtime/models/smollm3/triattention_kernels.h" +#endif + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace trtmc { + +namespace { + +using json = nlohmann::json; + +constexpr float kMaskedScore = -1.0e4F; +constexpr float kEps = 1.0e-6F; +constexpr float kAbsFloor = 1.0e-8F; + +void require_ta(bool cond, const char* msg) { + if (!cond) + throw std::runtime_error(msg); +} + +Smollm3TriAttentionScoreAggregation parse_score_aggregation(const std::string& value) { + if (value == "mean") + return Smollm3TriAttentionScoreAggregation::kMean; + if (value == "max") + return Smollm3TriAttentionScoreAggregation::kMax; + throw std::runtime_error("Unsupported TriAttention score_aggregation: " + value); +} + +const char* score_aggregation_name(Smollm3TriAttentionScoreAggregation value) { + switch (value) { + case Smollm3TriAttentionScoreAggregation::kMean: + return "mean"; + case Smollm3TriAttentionScoreAggregation::kMax: + return "max"; + } + return "unknown"; +} + +Smollm3TriAttentionRopeStyle parse_rope_style(const std::string& value) { + if (value == "interleaved") + return Smollm3TriAttentionRopeStyle::kInterleaved; + if (value.empty() || value == "half") + return Smollm3TriAttentionRopeStyle::kHalf; + throw std::runtime_error("Unsupported TriAttention rope_style: " + value); +} + +std::vector make_offsets(int32_t max_length) { + std::vector out; + for (int32_t current = 1; current > 0 && current <= max_length; current *= 2) + out.push_back(static_cast(current)); + return out; +} + +float fp16_to_float(uint16_t bits) { + const uint32_t sign = static_cast(bits & 0x8000U) << 16; + int32_t exponent = static_cast((bits >> 10) & 0x1FU); + uint32_t mantissa = bits & 0x03FFU; + + uint32_t out_bits = 0; + if (exponent == 0) { + if (mantissa == 0) { + out_bits = sign; + } else { + exponent = 1; + while ((mantissa & 0x0400U) == 0U) { + mantissa <<= 1; + --exponent; + } + mantissa &= 0x03FFU; + out_bits = + sign | (static_cast(exponent + (127 - 15)) << 23) | (mantissa << 13); + } + } else if (exponent == 0x1FU) { + out_bits = sign | 0x7F800000U | (mantissa << 13); + } else { + out_bits = sign | (static_cast(exponent + (127 - 15)) << 23) | (mantissa << 13); + } + + float out = 0.0F; + std::memcpy(&out, &out_bits, sizeof(out)); + return out; +} + +float bf16_to_float(uint16_t bits) { + const uint32_t out_bits = static_cast(bits) << 16; + float out = 0.0F; + std::memcpy(&out, &out_bits, sizeof(out)); + return out; +} + +std::vector parse_float_array(const json& array_value, const char* label) { + if (!array_value.is_array()) + throw std::runtime_error(std::string("TriAttention ") + label + " must be an array"); + std::vector out; + out.reserve(array_value.size()); + for (const auto& item : array_value) + out.push_back(item.get()); + return out; +} + +std::vector> parse_float_matrix(const json& array_value, const char* label) { + if (!array_value.is_array()) { + throw std::runtime_error(std::string("TriAttention ") + label + " must be a 2D array"); + } + std::vector> out; + out.reserve(array_value.size()); + for (const auto& row : array_value) + out.push_back(parse_float_array(row, label)); + return out; +} + +void maybe_generate_default_names(int32_t num_layers, Smollm3KvCacheNames& names) { + if (!names.cache_k.empty()) + return; + names.cache_k.reserve(static_cast(num_layers)); + names.cache_v.reserve(static_cast(num_layers)); + names.present_k.reserve(static_cast(num_layers)); + names.present_v.reserve(static_cast(num_layers)); + for (int32_t i = 0; i < num_layers; ++i) { + std::string suffix = "_" + std::to_string(i); + names.cache_k.push_back("cache_k" + suffix); + names.cache_v.push_back("cache_v" + suffix); + names.present_k.push_back("present_k" + suffix); + names.present_v.push_back("present_v" + suffix); + } +} + +std::string sampled_head_key(int32_t layer, int32_t head) { + char key[32]; + std::snprintf(key, sizeof(key), "layer%02d_head%02d", layer, head); + return std::string(key); +} + +float complex_abs(float real, float imag) { + return std::sqrt(std::max(real * real + imag * imag, kAbsFloor)); +} + +// --- Registry-backed value reader ------------------------------------------ +// +// apply_layer_value_ overlays a value from the runtime config onto an +// out-param IFF the registry has a non-default layer value for the field +// (i.e. something above SchemaDefault contributed). When the source is +// SchemaDefault we leave the out-param alone so legacy bundle-JSON values +// keep precedence for fields the caller never touched from the CLI. +// +// Previously this file contained a cluster of std::getenv readers +// (triattention_debug_enabled, triattention_profile_enabled, etc.) plus +// override helpers that patched Smollm3TriAttentionConfig with +// TRTMC_TRIATTN_OVERRIDE_* values. All of them are deleted here — values +// now flow through the config registry exclusively. +template +bool registry_has_value(const ::trtmc::config::ConfigBundle& bundle, const std::string& field) { + try { + return bundle.source_of("triattention", field) != ::trtmc::config::Layer::SchemaDefault; + } catch (const std::exception&) { + return false; + } +} + +template +void apply_layer_value(const ::trtmc::config::ConfigBundle& bundle, const std::string& field, + T& out) { + if (!registry_has_value(bundle, field)) + return; + try { + out = bundle.get("triattention", field); + } catch (const std::exception&) { /* type mismatch — skip */ + } +} + +void apply_aggregation_from_registry(const ::trtmc::config::ConfigBundle& bundle, + const std::string& field, + Smollm3TriAttentionScoreAggregation& out) { + if (!registry_has_value(bundle, field)) + return; + try { + out = parse_score_aggregation(bundle.get("triattention", field)); + } catch (const std::exception&) { /* keep previous value */ + } +} + +using Clock = std::chrono::steady_clock; + +double elapsed_ms(const Clock::time_point start) { + return std::chrono::duration(Clock::now() - start).count(); +} + +float score_full_stddev_or_one(const std::vector& scores, float mean) { + if (scores.size() <= 1U) + return 1.0F; + float var = 0.0F; + for (const float value : scores) { + const float delta = value - mean; + var += delta * delta; + } + const float denom = static_cast(scores.size() - 1U); + const float stddev = std::sqrt(std::max(var / denom, 0.0F)); + return stddev < kEps ? 1.0F : stddev; +} + +int32_t round_up_rows(int32_t value, int32_t bucket, int32_t maximum) { + if (bucket <= 1) + return std::min(std::max(value, 1), maximum); + const int32_t rounded = ((std::max(value, 1) + bucket - 1) / bucket) * bucket; + return std::min(rounded, maximum); +} + +} // namespace + +namespace { + +// Overlay core-runtime fields from the registry on top of whatever the +// legacy JSON path produced. Session > Platform > BundleDefault > BuildTime +// win; SchemaDefault reads are skipped so JSON-only bundles keep their +// values. +void overlay_core_runtime_from_registry(Smollm3TriAttentionConfig& cfg, + const ::trtmc::config::ConfigBundle& bundle) { + apply_layer_value(bundle, "enabled", cfg.enabled); + apply_layer_value(bundle, "kv_budget", cfg.kv_budget); + apply_layer_value(bundle, "divide_length", cfg.divide_length); + apply_layer_value(bundle, "recent_window", cfg.recent_window); + apply_aggregation_from_registry(bundle, "score_aggregation", cfg.score_aggregation); + apply_aggregation_from_registry(bundle, "per_layer_aggregation", cfg.per_layer_aggregation); + apply_layer_value(bundle, "count_prompt_tokens", cfg.count_prompt_tokens); + apply_layer_value(bundle, "protect_prefill", cfg.protect_prefill); + apply_layer_value(bundle, "disable_mlr", cfg.disable_mlr); + apply_layer_value(bundle, "disable_trig", cfg.disable_trig); + apply_layer_value(bundle, "stats_section", cfg.stats_section); + apply_layer_value(bundle, "offset_max_length", cfg.offset_max_length); +} + +// Populate the debug/profile fields from the registry. These have no +// legacy JSON representation — they previously came from +// TRTMC_TRIATTN_* env vars, which are now gone. +void fill_debug_from_registry(Smollm3TriAttentionConfig& cfg, + const ::trtmc::config::ConfigBundle& bundle) { + apply_layer_value(bundle, "debug", cfg.debug); + apply_layer_value(bundle, "profile", cfg.profile); + apply_layer_value(bundle, "runtime_bucket_rows", cfg.runtime_bucket_rows); + apply_layer_value(bundle, "disable_gpu_selection", cfg.disable_gpu_selection); + apply_layer_value(bundle, "disable_gpu_compaction", cfg.disable_gpu_compaction); + apply_layer_value(bundle, "disable_gpu_state", cfg.disable_gpu_state); + apply_layer_value(bundle, "zero_tail", cfg.zero_tail); + apply_layer_value(bundle, "dump_keep_path", cfg.dump_keep_path); + apply_layer_value(bundle, "dump_compaction_index", cfg.dump_compaction_index); + apply_layer_value(bundle, "abort_after_dump", cfg.abort_after_dump); + apply_layer_value(bundle, "dump_score_cache", cfg.dump_score_cache); + apply_layer_value(bundle, "dump_score_values", cfg.dump_score_values); +} + +static void fill_core_from_legacy_json(Smollm3TriAttentionConfig& cfg, + const std::string& config_json, int32_t max_cache_length) { + if (config_json.find("\"triattention\"") == std::string::npos) + return; + const json root = json::parse(config_json); + const auto it = root.find("triattention"); + if (it == root.end() || !it->is_object()) + return; + cfg.enabled = it->value("enabled", false); + cfg.kv_budget = it->value("kv_budget", max_cache_length); + cfg.divide_length = it->value("divide_length", 128); + cfg.recent_window = it->value("recent_window", 128); + cfg.score_aggregation = + parse_score_aggregation(it->value("score_aggregation", std::string("mean"))); + cfg.per_layer_aggregation = + parse_score_aggregation(it->value("per_layer_aggregation", std::string("mean"))); + cfg.count_prompt_tokens = it->value("count_prompt_tokens", true); + cfg.protect_prefill = it->value("protect_prefill", true); + cfg.disable_mlr = it->value("disable_mlr", false); + cfg.disable_trig = it->value("disable_trig", false); + cfg.stats_section = it->value("stats_section", std::string("triattention_stats.json")); + cfg.offset_max_length = it->value("offset_max_length", 65536); +} + +static void validate_triattention_config(const Smollm3TriAttentionConfig& cfg, + int32_t max_cache_length) { + if (cfg.kv_budget < 1) + throw std::runtime_error("TriAttention kv_budget must be >= 1"); + if (cfg.kv_budget > max_cache_length) + throw std::runtime_error("TriAttention kv_budget cannot exceed engine max_cache_length"); + if (cfg.divide_length < 1) + throw std::runtime_error("TriAttention divide_length must be >= 1"); + if (cfg.recent_window < 0) + throw std::runtime_error("TriAttention recent_window must be >= 0"); + if (cfg.offset_max_length < 1) + throw std::runtime_error("TriAttention offset_max_length must be >= 1"); +} + +} // namespace + +Smollm3TriAttentionConfig +smollm3_parse_triattention_bundle_config(const std::string& config_json, int32_t max_cache_length, + const ::trtmc::config::ConfigBundle* runtime_config) { + Smollm3TriAttentionConfig cfg; + // Legacy bundle path: pull core fields from the root-level + // "triattention" object. New bundles route the same values through + // the generic `defaults:` block which becomes the BundleDefault layer + // in runtime_config. + fill_core_from_legacy_json(cfg, config_json, max_cache_length); + // Overlay registry-supplied values. Session/platform/bundle-default + // layers win; SchemaDefault reads are skipped. Debug fields come + // exclusively from the registry (no legacy JSON path). + if (runtime_config != nullptr) { + overlay_core_runtime_from_registry(cfg, *runtime_config); + fill_debug_from_registry(cfg, *runtime_config); + } + if (!cfg.enabled) + return cfg; + validate_triattention_config(cfg, max_cache_length); + return cfg; +} + +namespace { + +void fill_stats_header(Smollm3TriAttentionStats& stats, const json& root, + int32_t num_attention_heads, int32_t num_key_value_heads, + int32_t num_layers) { + stats.head_dim = root.value("head_dim", 0); + stats.rope_style = parse_rope_style(root.value("rope_style", std::string("half"))); + stats.rope_theta = root.value("rope_theta", 10000.0F); + stats.num_attention_heads = std::max(root.value("num_attention_heads", num_attention_heads), 1); + stats.num_key_value_heads = std::max(root.value("num_key_value_heads", num_key_value_heads), 1); + stats.stats_head_count = std::max(root.value("stats_head_count", 0), 0); + stats.num_layers = std::max(root.value("num_layers", num_layers), 1); + require_ta(stats.head_dim > 0 && (stats.head_dim % 2) == 0, + "TriAttention stats head_dim must be a positive even number"); +} + +void fill_stats_inv_freq(Smollm3TriAttentionStats& stats, const json& root) { + if (root.contains("inv_freq")) { + stats.inv_freq = parse_float_array(root["inv_freq"], "inv_freq"); + } else { + const int32_t half_dim = stats.head_dim / 2; + stats.inv_freq.reserve(static_cast(half_dim)); + for (int32_t i = 0; i < half_dim; ++i) { + const float exponent = + (2.0F * static_cast(i)) / static_cast(stats.head_dim); + stats.inv_freq.push_back(1.0F / std::pow(stats.rope_theta, exponent)); + } + } + require_ta(static_cast(stats.inv_freq.size()) == stats.head_dim / 2, + "TriAttention inv_freq size does not match head_dim / 2"); +} + +void append_sampled_heads_from_array(Smollm3TriAttentionStats& stats, const json& sampled_root, + int32_t head_upper_bound) { + for (const auto& item : sampled_root) { + require_ta(item.is_array() && item.size() == 2, + "TriAttention sampled_heads entries must be [layer, head]"); + const int32_t layer = item[0].get(); + const int32_t head = item[1].get(); + require_ta(layer >= 0 && layer < stats.num_layers, + "TriAttention sampled head layer is out of range"); + require_ta(head >= 0 && head < head_upper_bound, + "TriAttention sampled head index is out of range"); + stats.sampled_score_heads_by_layer[static_cast(layer)].push_back(head); + } +} + +void populate_sampled_heads(Smollm3TriAttentionStats& stats, const json& root, + int32_t head_upper_bound) { + stats.sampled_score_heads_by_layer.assign(static_cast(stats.num_layers), {}); + if (root.contains("sampled_heads") && root["sampled_heads"].is_array()) + append_sampled_heads_from_array(stats, root["sampled_heads"], head_upper_bound); + + bool any_sampled = false; + for (auto& heads : stats.sampled_score_heads_by_layer) { + std::sort(heads.begin(), heads.end()); + heads.erase(std::unique(heads.begin(), heads.end()), heads.end()); + any_sampled = any_sampled || !heads.empty(); + } + if (any_sampled) + return; + for (auto& heads : stats.sampled_score_heads_by_layer) { + heads.resize(static_cast(head_upper_bound)); + std::iota(heads.begin(), heads.end(), 0); + } +} + +int32_t infer_sampled_head_upper_bound(const json& root) { + if (!root.contains("sampled_heads") || !root["sampled_heads"].is_array()) + return 0; + int32_t upper_bound = 0; + for (const auto& item : root["sampled_heads"]) { + if (!item.is_array() || item.size() != 2) + continue; + upper_bound = std::max(upper_bound, item[1].get() + 1); + } + return upper_bound; +} + +void copy_dense_head_row(Smollm3TriAttentionHeadStats& dst, + const std::vector>& q_mean_real, + const std::vector>& q_mean_imag, + const std::vector>& q_abs_mean, + const std::vector>& freq_scale_sq, int32_t score_head, + int32_t half_dim) { + const std::size_t src_idx = static_cast(score_head); + require_ta(static_cast(q_mean_real[src_idx].size()) == half_dim && + static_cast(q_mean_imag[src_idx].size()) == half_dim && + static_cast(q_abs_mean[src_idx].size()) == half_dim && + static_cast(freq_scale_sq[src_idx].size()) == half_dim, + "TriAttention layer_stats frequency count does not match head_dim / 2"); + const auto base = src_idx * static_cast(half_dim); + for (int32_t d = 0; d < half_dim; ++d) { + const auto dst_idx = base + static_cast(d); + dst.q_mean_real[dst_idx] = q_mean_real[src_idx][static_cast(d)]; + dst.q_mean_imag[dst_idx] = q_mean_imag[src_idx][static_cast(d)]; + dst.q_abs_mean[dst_idx] = q_abs_mean[src_idx][static_cast(d)]; + dst.freq_scale_sq[dst_idx] = freq_scale_sq[src_idx][static_cast(d)]; + } +} + +bool fill_layer_stats_from_dense(Smollm3TriAttentionHeadStats& dst, const json& layer_node, + int32_t& inferred_stats_heads, int32_t half_dim) { + const auto q_mean_real = + parse_float_matrix(layer_node["q_mean_real"], "layer_stats.q_mean_real"); + const auto q_mean_imag = + parse_float_matrix(layer_node["q_mean_imag"], "layer_stats.q_mean_imag"); + const auto q_abs_mean = parse_float_matrix(layer_node["q_abs_mean"], "layer_stats.q_abs_mean"); + const int32_t row_count = static_cast(q_mean_real.size()); + if (inferred_stats_heads <= 0) + inferred_stats_heads = row_count; + std::vector> freq_scale_sq( + static_cast(inferred_stats_heads), + std::vector(static_cast(half_dim), 1.0F)); + if (layer_node.contains("freq_scale_sq")) + freq_scale_sq = + parse_float_matrix(layer_node["freq_scale_sq"], "layer_stats.freq_scale_sq"); + require_ta(row_count == inferred_stats_heads && + static_cast(q_mean_imag.size()) == inferred_stats_heads && + static_cast(q_abs_mean.size()) == inferred_stats_heads && + static_cast(freq_scale_sq.size()) == inferred_stats_heads, + "TriAttention layer_stats head count is inconsistent"); + const auto flat_size = + static_cast(inferred_stats_heads) * static_cast(half_dim); + dst.q_mean_real.resize(flat_size); + dst.q_mean_imag.resize(flat_size); + dst.q_abs_mean.resize(flat_size); + dst.freq_scale_sq.resize(flat_size); + for (int32_t score_head = 0; score_head < inferred_stats_heads; ++score_head) + copy_dense_head_row(dst, q_mean_real, q_mean_imag, q_abs_mean, freq_scale_sq, score_head, + half_dim); + return true; +} + +bool try_parse_dense_layer_stats(Smollm3TriAttentionStats& stats, const json& root, + int32_t half_dim) { + if (!root.contains("layer_stats") || !root["layer_stats"].is_object() || + root["layer_stats"].empty()) + return false; + stats.layer_stats.resize(static_cast(stats.num_layers)); + int32_t inferred_stats_heads = stats.stats_head_count; + bool any_stats = false; + for (int32_t layer = 0; layer < stats.num_layers; ++layer) { + auto layer_it = root["layer_stats"].find(std::to_string(layer)); + if (layer_it == root["layer_stats"].end() || !layer_it->is_object()) + continue; + fill_layer_stats_from_dense(stats.layer_stats[static_cast(layer)], *layer_it, + inferred_stats_heads, half_dim); + any_stats = true; + } + if (!any_stats) + return false; + stats.stats_head_count = std::max(inferred_stats_heads, 1); + populate_sampled_heads(stats, root, stats.stats_head_count); + return true; +} + +void init_layer_stats_empty(Smollm3TriAttentionStats& stats, int32_t half_dim) { + stats.layer_stats.resize(static_cast(stats.num_layers)); + const auto flat_size = + static_cast(stats.stats_head_count) * static_cast(half_dim); + for (auto& layer : stats.layer_stats) { + layer.q_mean_real.assign(flat_size, 0.0F); + layer.q_mean_imag.assign(flat_size, 0.0F); + layer.q_abs_mean.assign(flat_size, 0.0F); + layer.freq_scale_sq.assign(flat_size, 1.0F); + } +} + +void accumulate_sparse_entry(Smollm3TriAttentionStats& stats, const json& raw_stats, int32_t layer, + int32_t head, int32_t half_dim, std::vector& group_counts) { + const std::string key = sampled_head_key(layer, head); + auto stats_it = raw_stats.find(key); + require_ta(stats_it != raw_stats.end() && stats_it->is_object(), + "TriAttention stats payload is missing entry"); + const auto q_mean_real = parse_float_array((*stats_it)["q_mean_real"], "q_mean_real"); + const auto q_mean_imag = parse_float_array((*stats_it)["q_mean_imag"], "q_mean_imag"); + const auto q_abs_mean = parse_float_array((*stats_it)["q_abs_mean"], "q_abs_mean"); + require_ta(static_cast(q_mean_real.size()) == half_dim && + static_cast(q_mean_imag.size()) == half_dim && + static_cast(q_abs_mean.size()) == half_dim, + "TriAttention sparse stats entry does not match head_dim / 2"); + auto& layer_stats = stats.layer_stats[static_cast(layer)]; + const auto base = static_cast(head) * static_cast(half_dim); + for (int32_t d = 0; d < half_dim; ++d) { + const auto idx = base + static_cast(d); + layer_stats.q_mean_real[idx] += q_mean_real[static_cast(d)]; + layer_stats.q_mean_imag[idx] += q_mean_imag[static_cast(d)]; + layer_stats.q_abs_mean[idx] += q_abs_mean[static_cast(d)]; + } + ++group_counts[static_cast(layer * stats.stats_head_count + head)]; +} + +bool finalize_sparse_stats(Smollm3TriAttentionStats& stats, + const std::vector& group_counts, int32_t half_dim) { + bool any_stats = false; + for (int32_t layer = 0; layer < stats.num_layers; ++layer) { + auto& layer_stats = stats.layer_stats[static_cast(layer)]; + for (int32_t score_head = 0; score_head < stats.stats_head_count; ++score_head) { + const int32_t count = + group_counts[static_cast(layer * stats.stats_head_count + score_head)]; + if (count <= 0) + continue; + any_stats = true; + const auto base = + static_cast(score_head) * static_cast(half_dim); + for (int32_t d = 0; d < half_dim; ++d) { + const auto idx = base + static_cast(d); + layer_stats.q_mean_real[idx] /= static_cast(count); + layer_stats.q_mean_imag[idx] /= static_cast(count); + layer_stats.q_abs_mean[idx] /= static_cast(count); + } + } + } + return any_stats; +} + +void parse_sparse_stats(Smollm3TriAttentionStats& stats, const json& root, int32_t half_dim) { + require_ta(root.contains("sampled_heads") && root["sampled_heads"].is_array(), + "TriAttention stats payload is missing sampled_heads"); + require_ta(root.contains("stats") && root["stats"].is_object(), + "TriAttention stats payload is missing stats object"); + require_ta(stats.num_attention_heads % stats.num_key_value_heads == 0, + "TriAttention num_attention_heads must be divisible by num_key_value_heads"); + if (stats.stats_head_count <= 0) { + stats.stats_head_count = + std::max({infer_sampled_head_upper_bound(root), stats.num_key_value_heads, 1}); + } + populate_sampled_heads(stats, root, stats.stats_head_count); + init_layer_stats_empty(stats, half_dim); + + std::vector group_counts( + static_cast(stats.num_layers * stats.stats_head_count), 0); + const auto& raw_stats = root["stats"]; + for (const auto& item : root["sampled_heads"]) { + require_ta(item.is_array() && item.size() == 2, + "TriAttention sampled_heads entries must be [layer, head]"); + const int32_t layer = item.at(0).get(); + const int32_t head = item.at(1).get(); + if (layer < 0 || layer >= stats.num_layers) + continue; + if (head < 0 || head >= stats.stats_head_count) + continue; + accumulate_sparse_entry(stats, raw_stats, layer, head, half_dim, group_counts); + } + require_ta(finalize_sparse_stats(stats, group_counts, half_dim), + "TriAttention stats payload has no usable sampled heads"); +} + +} // namespace + +Smollm3TriAttentionStats smollm3_parse_triattention_stats_json(const std::string& stats_json, + int32_t num_attention_heads, + int32_t num_key_value_heads, + int32_t num_layers) { + Smollm3TriAttentionStats stats; + const json root = json::parse(stats_json); + fill_stats_header(stats, root, num_attention_heads, num_key_value_heads, num_layers); + fill_stats_inv_freq(stats, root); + const int32_t half_dim = stats.head_dim / 2; + if (try_parse_dense_layer_stats(stats, root, half_dim)) + return stats; + parse_sparse_stats(stats, root, half_dim); + return stats; +} + +void Smollm3TriAttentionKvCache::validate_shapes() { + require_ta(config_.kv_budget >= 1 && config_.kv_budget <= max_length_, + "TriAttention kv_budget must be within [1, max_length]"); + require_ta(stats_.head_dim > 0 && (stats_.head_dim % 2) == 0, + "TriAttention stats head_dim must be a positive even number"); + require_ta(num_kv_heads_ > 0, "TriAttention num_kv_heads must be positive"); + query_head_count_ = kv_dim_ / stats_.head_dim; + require_ta(query_head_count_ > 0 && (kv_dim_ % stats_.head_dim) == 0, + "TriAttention kv_dim must be divisible by head_dim"); + require_ta((query_head_count_ % num_kv_heads_) == 0, + "TriAttention expanded cache heads must be divisible by kv heads"); + require_ta(stats_.num_attention_heads > 0 && stats_.num_key_value_heads > 0 && + (stats_.num_attention_heads % stats_.num_key_value_heads) == 0, + "TriAttention stats attention head count must be divisible by kv heads"); + query_group_size_ = query_head_count_ / num_kv_heads_; + cache_head_count_ = num_kv_heads_; + require_ta(stats_.stats_head_count > 0 && (stats_.stats_head_count % cache_head_count_) == 0, + "TriAttention stats_head_count must be divisible by runtime kv heads"); + score_group_size_ = stats_.stats_head_count / cache_head_count_; + require_ta(score_group_size_ > 0, "TriAttention score_group_size must be positive"); +} + +void Smollm3TriAttentionKvCache::normalize_sampled_heads() { + if (stats_.sampled_score_heads_by_layer.size() != static_cast(num_layers_)) + stats_.sampled_score_heads_by_layer.assign(static_cast(num_layers_), {}); + + const int32_t dense_head_upper_bound = stats_.stats_head_count; + bool any_sampled_heads = false; + for (auto& heads : stats_.sampled_score_heads_by_layer) { + std::sort(heads.begin(), heads.end()); + heads.erase(std::unique(heads.begin(), heads.end()), heads.end()); + for (const int32_t head : heads) { + if (head < 0 || head >= dense_head_upper_bound) + throw std::runtime_error("TriAttention sampled score head index is out of range"); + } + any_sampled_heads = any_sampled_heads || !heads.empty(); + } + if (!any_sampled_heads) { + for (auto& heads : stats_.sampled_score_heads_by_layer) { + heads.resize(static_cast(dense_head_upper_bound)); + std::iota(heads.begin(), heads.end(), 0); + } + } +} + +void Smollm3TriAttentionKvCache::log_init_debug() const { + if (!config_.debug) + return; + std::cerr << "[trtmc.triattention] init kv_budget=" << config_.kv_budget + << " divide_length=" << config_.divide_length + << " recent_window=" << config_.recent_window + << " per_layer_aggregation=" << score_aggregation_name(config_.per_layer_aggregation) + << " count_prompt_tokens=" << (config_.count_prompt_tokens ? 1 : 0) + << " protect_prefill=" << (config_.protect_prefill ? 1 : 0) + << " disable_mlr=" << (config_.disable_mlr ? 1 : 0) + << " disable_trig=" << (config_.disable_trig ? 1 : 0) << " kv_heads=" << num_kv_heads_ + << " query_heads=" << query_head_count_ << " cache_heads=" << cache_head_count_ + << " score_group=" << score_group_size_ + << " layers_with_stats=" << stats_.layer_stats.size() << '\n'; +} + +void Smollm3TriAttentionKvCache::allocate_layer_tensors() { + cache_k_.reserve(static_cast(num_layers_)); + cache_v_.reserve(static_cast(num_layers_)); + present_k_.reserve(static_cast(num_layers_)); + present_v_.reserve(static_cast(num_layers_)); + for (int32_t i = 0; i < num_layers_; ++i) { + cache_k_.emplace_back(std::vector{max_length_, kv_dim_}, cache_dtype_, stream_); + cache_v_.emplace_back(std::vector{max_length_, kv_dim_}, cache_dtype_, stream_); + present_k_.emplace_back(std::vector{1, kv_dim_}, cache_dtype_, stream_); + present_v_.emplace_back(std::vector{1, kv_dim_}, cache_dtype_, stream_); + } + mask_buf_.resize(static_cast(max_length_) + 1U); + cache_positions_.reserve(static_cast(max_length_)); + cache_positions_per_head_.resize(static_cast(cache_head_count_)); + for (auto& head_positions : cache_positions_per_head_) + head_positions.reserve(static_cast(max_length_)); +} + +Smollm3TriAttentionKvCache::Smollm3TriAttentionKvCache(int32_t num_layers, int32_t num_kv_heads, + int32_t max_length, int32_t kv_dim, + cudaStream_t stream, + Smollm3TriAttentionConfig config, + Smollm3TriAttentionStats stats, + DType cache_dtype, Smollm3KvCacheNames names) + : num_layers_(num_layers), num_kv_heads_(num_kv_heads), max_length_(max_length), + kv_dim_(kv_dim), stream_(stream), cache_dtype_(cache_dtype), + cache_element_size_(dtype_size(cache_dtype)), names_(std::move(names)), + config_(std::move(config)), stats_(std::move(stats)), + offsets_(make_offsets(config_.offset_max_length)) { + validate_shapes(); + normalize_sampled_heads(); + profile_enabled_ = config_.profile; + log_init_debug(); + maybe_generate_default_names(num_layers_, names_); + allocate_layer_tensors(); +#ifdef TRTMC_HAS_CUDA_KERNELS + if (!config_.disable_gpu_state) + initialize_gpu_state(); +#endif + reset(); +} + +#ifdef TRTMC_HAS_CUDA_KERNELS +void Smollm3TriAttentionKvCache::allocate_core_selection_buffers(int32_t half_dim) { + candidate_indices_device_ = DeviceTensor({max_length_}, DType::kInt32, stream_); + keep_indices_device_ = DeviceTensor({static_cast(cache_head_count_) * max_length_}, + DType::kInt32, stream_); + positions_device_ = DeviceTensor({static_cast(cache_head_count_) * max_length_}, + DType::kInt32, stream_); + inv_freq_device_ = + DeviceTensor({static_cast(stats_.inv_freq.size())}, DType::kFloat32, stream_); + cos_phase_device_ = + DeviceTensor({static_cast(offsets_.size() * static_cast(half_dim))}, + DType::kFloat32, stream_); + sin_phase_device_ = + DeviceTensor({static_cast(offsets_.size() * static_cast(half_dim))}, + DType::kFloat32, stream_); + scratch_k_device_ = DeviceTensor({max_length_, kv_dim_}, cache_dtype_, stream_); + scratch_v_device_ = DeviceTensor({max_length_, kv_dim_}, cache_dtype_, stream_); + if (inv_freq_device_.ok()) + inv_freq_device_.copy_from_host(stats_.inv_freq.data()); +} + +void Smollm3TriAttentionKvCache::build_layer_gpu_stats(int32_t layer, int32_t half_dim) { + const auto& host = stats_.layer_stats[static_cast(layer)]; + const auto& sampled_heads = + stats_.sampled_score_heads_by_layer[static_cast(layer)]; + auto& gpu = layer_gpu_stats_[static_cast(layer)]; + gpu.score_head_count = static_cast(sampled_heads.size()); + gpu.host_cache_head_indices.clear(); + if (host.q_mean_real.empty() || host.q_mean_imag.empty() || host.q_abs_mean.empty()) + return; + if (gpu.score_head_count <= 0) + return; + const auto expected_stats_size = + static_cast(stats_.stats_head_count) * static_cast(half_dim); + if (host.q_mean_real.size() != expected_stats_size || + host.q_mean_imag.size() != expected_stats_size || + host.q_abs_mean.size() != expected_stats_size || + host.freq_scale_sq.size() != expected_stats_size) + return; + + const auto n_heads = static_cast(gpu.score_head_count); + const auto flat = n_heads * static_cast(half_dim); + std::vector head_offsets(n_heads); + std::vector head_cache_indices(n_heads); + std::vector q_mean_real(flat), q_mean_imag(flat), q_abs_mean(flat), freq_scale_sq(flat); + gpu.host_cache_head_indices.resize(n_heads); + + for (int32_t sampled_idx = 0; sampled_idx < gpu.score_head_count; ++sampled_idx) { + const int32_t score_head = sampled_heads[static_cast(sampled_idx)]; + const int32_t cache_head = std::min(cache_head_count_ - 1, score_head / score_group_size_); + head_cache_indices[static_cast(sampled_idx)] = cache_head; + gpu.host_cache_head_indices[static_cast(sampled_idx)] = cache_head; + head_offsets[static_cast(sampled_idx)] = + cache_head * query_group_size_ * stats_.head_dim; + const auto src_base = + static_cast(score_head) * static_cast(half_dim); + const auto dst_base = + static_cast(sampled_idx) * static_cast(half_dim); + std::copy_n(host.q_mean_real.begin() + static_cast(src_base), half_dim, + q_mean_real.begin() + static_cast(dst_base)); + std::copy_n(host.q_mean_imag.begin() + static_cast(src_base), half_dim, + q_mean_imag.begin() + static_cast(dst_base)); + std::copy_n(host.q_abs_mean.begin() + static_cast(src_base), half_dim, + q_abs_mean.begin() + static_cast(dst_base)); + std::copy_n(host.freq_scale_sq.begin() + static_cast(src_base), half_dim, + freq_scale_sq.begin() + static_cast(dst_base)); + } + + gpu.head_offsets = DeviceTensor({gpu.score_head_count}, DType::kInt32, stream_); + gpu.head_cache_indices = DeviceTensor({gpu.score_head_count}, DType::kInt32, stream_); + gpu.q_mean_real = DeviceTensor({static_cast(gpu.score_head_count) * half_dim}, + DType::kFloat32, stream_); + gpu.q_mean_imag = DeviceTensor({static_cast(gpu.score_head_count) * half_dim}, + DType::kFloat32, stream_); + gpu.q_abs_mean = DeviceTensor({static_cast(gpu.score_head_count) * half_dim}, + DType::kFloat32, stream_); + gpu.freq_scale_sq = DeviceTensor({static_cast(gpu.score_head_count) * half_dim}, + DType::kFloat32, stream_); + gpu.scores = DeviceTensor({static_cast(gpu.score_head_count) * max_length_}, + DType::kFloat32, stream_); + + gpu.head_offsets.copy_from_host(head_offsets.data()); + gpu.head_cache_indices.copy_from_host(head_cache_indices.data()); + gpu.q_mean_real.copy_from_host(q_mean_real.data()); + gpu.q_mean_imag.copy_from_host(q_mean_imag.data()); + gpu.q_abs_mean.copy_from_host(q_abs_mean.data()); + gpu.freq_scale_sq.copy_from_host(freq_scale_sq.data()); +} + +void Smollm3TriAttentionKvCache::initialize_gpu_state() { + const int32_t half_dim = stats_.head_dim / 2; + if (half_dim <= 0 || offsets_.empty() || stats_.layer_stats.empty() || query_head_count_ <= 0) + return; + allocate_core_selection_buffers(half_dim); + layer_gpu_stats_.resize(static_cast(num_layers_)); + for (int32_t layer = 0; layer < num_layers_; ++layer) { + if (layer >= static_cast(stats_.layer_stats.size())) + break; + build_layer_gpu_stats(layer, half_dim); + } + cudaStreamSynchronize(stream_); +} + +bool Smollm3TriAttentionKvCache::core_selection_buffers_ready() const { + return candidate_indices_device_.ok() && keep_indices_device_.ok() && positions_device_.ok() && + inv_freq_device_.ok() && cos_phase_device_.ok() && sin_phase_device_.ok() && + scratch_k_device_.ok() && scratch_v_device_.ok(); +} + +bool Smollm3TriAttentionKvCache::layer_gpu_stats_ready(const LayerGpuStats& layer) { + if (layer.score_head_count == 0) + return true; + return layer.head_offsets.ok() && layer.head_cache_indices.ok() && layer.q_mean_real.ok() && + layer.q_mean_imag.ok() && layer.q_abs_mean.ok() && layer.freq_scale_sq.ok() && + layer.scores.ok(); +} + +bool Smollm3TriAttentionKvCache::can_use_gpu_selection() const { + if (config_.disable_gpu_selection) + return false; + if (!core_selection_buffers_ready()) + return false; + for (const auto& layer : layer_gpu_stats_) { + if (!layer_gpu_stats_ready(layer)) + return false; + } + return true; +} +#endif + +void Smollm3TriAttentionKvCache::build_attention_mask(std::vector& mask) const { + const auto width = static_cast(max_length_) + 1U; + mask.assign(width, kMaskedScore); + for (int32_t i = 0; i < cache_length_; ++i) + mask[static_cast(i)] = 0.0F; + mask.back() = 0.0F; +} + +int32_t Smollm3TriAttentionKvCache::preferred_cache_rows() const { + if (!dynamic_binding_enabled_) + return max_length_; + const int32_t bucket_rows = config_.runtime_bucket_rows; + return round_up_rows(std::max(cache_length_, 1), bucket_rows, max_length_); +} + +void Smollm3TriAttentionKvCache::prepare_step(TensorMap& inputs, int32_t /*seq_len*/) { + if (has_position_input_) { + pos_buf_ = absolute_position_; + Tensor pos_t; + pos_t.data = &pos_buf_; + pos_t.shape = {1}; + pos_t.dtype = DType::kInt32; + inputs[names_.position_id] = pos_t; + } + + const int32_t cache_rows = dynamic_binding_enabled_ ? preferred_cache_rows() : max_length_; + const int32_t mask_width = dynamic_binding_enabled_ ? (cache_rows + 1) : (max_length_ + 1); + if (dynamic_binding_enabled_ && bound_module_ != nullptr && cache_rows != bound_cache_rows_) { + const std::vector cache_shape{cache_rows, kv_dim_}; + for (int32_t i = 0; i < num_layers_; ++i) { + const auto li = static_cast(i); + bound_module_->bind_external(names_.cache_k[li], cache_k_[li].data(), cache_shape); + bound_module_->bind_external(names_.cache_v[li], cache_v_[li].data(), cache_shape); + } + bound_cache_rows_ = cache_rows; + } + + std::fill(mask_buf_.begin(), mask_buf_.begin() + mask_width, kMaskedScore); + for (int32_t i = 0; i < cache_length_; ++i) + mask_buf_[static_cast(i)] = 0.0F; + mask_buf_[static_cast(mask_width - 1)] = 0.0F; + + Tensor mask_t; + mask_t.data = mask_buf_.data(); + mask_t.shape = dynamic_binding_enabled_ + ? std::vector{1, mask_width} + : std::vector{static_cast(mask_buf_.size())}; + mask_t.dtype = DType::kFloat32; + inputs[names_.attention_mask] = mask_t; +} + +void Smollm3TriAttentionKvCache::bind_to(TrtModule& module) { + bound_module_ = &module; + has_position_input_ = module.has_input(names_.position_id); + dynamic_binding_enabled_ = + !names_.cache_k.empty() && module.input_is_dynamic(names_.cache_k.front()); + bound_cache_rows_ = 0; + const int32_t initial_cache_rows = + dynamic_binding_enabled_ ? preferred_cache_rows() : max_length_; + const std::vector cache_shape{initial_cache_rows, kv_dim_}; + + for (int32_t i = 0; i < num_layers_; ++i) { + const auto li = static_cast(i); + if (dynamic_binding_enabled_) { + module.bind_external(names_.cache_k[li], cache_k_[li].data(), cache_shape); + module.bind_external(names_.cache_v[li], cache_v_[li].data(), cache_shape); + bound_cache_rows_ = initial_cache_rows; + } else { + module.bind_external(names_.cache_k[li], cache_k_[li].data()); + module.bind_external(names_.cache_v[li], cache_v_[li].data()); + } + module.bind_external(names_.present_k[li], present_k_[li].data()); + module.bind_external(names_.present_v[li], present_v_[li].data()); + } +} + +std::vector Smollm3TriAttentionKvCache::copy_cache_rows_to_host( + const DeviceTensor& tensor, int32_t rows, Smollm3TriAttentionCompactionProfile* profile) const { + const auto count = static_cast(rows) * static_cast(kv_dim_); + std::vector out(count, 0.0F); + if (rows <= 0) + return out; + + const auto raw_bytes = count * cache_element_size_; + std::vector raw(raw_bytes); + const auto copy_start = Clock::now(); + cudaStreamSynchronize(stream_); + cudaMemcpy(raw.data(), tensor.data(), raw_bytes, cudaMemcpyDeviceToHost); + if (profile != nullptr) { + profile->host_copy_ms += elapsed_ms(copy_start); + profile->host_copy_bytes += raw_bytes; + } + + if (cache_dtype_ == DType::kFloat32) { + std::memcpy(out.data(), raw.data(), raw_bytes); + return out; + } + + const auto convert_start = Clock::now(); + const auto* raw_u16 = reinterpret_cast(raw.data()); + for (std::size_t i = 0; i < count; ++i) { + if (cache_dtype_ == DType::kFloat16) + out[i] = fp16_to_float(raw_u16[i]); + else + out[i] = bf16_to_float(raw_u16[i]); + } + if (profile != nullptr) + profile->host_convert_ms += elapsed_ms(convert_start); + return out; +} + +void Smollm3TriAttentionKvCache::sync_shared_positions_from_head0() { + if (cache_positions_per_head_.empty()) { + cache_positions_.clear(); + return; + } + cache_positions_ = cache_positions_per_head_.front(); +} + +int32_t Smollm3TriAttentionKvCache::count_prefix_rows() const { + const int32_t prefix_limit = + prompt_end_position_ > 0 ? prompt_end_position_ : planned_prompt_length_; + if (prefix_limit <= 0) + return 0; + int32_t count = 0; + for (int32_t pos : cache_positions_) { + if (pos < prefix_limit) + ++count; + } + return count; +} + +std::vector Smollm3TriAttentionKvCache::build_reserve_mask(int32_t total_tokens, + int32_t old_budget) const { + std::vector reserve_mask(static_cast(total_tokens), 0); + const int32_t reserve_recent = + std::min({std::max(config_.recent_window, 0), total_tokens, old_budget}); + if (reserve_recent > 0) { + for (int32_t i = total_tokens - reserve_recent; i < total_tokens; ++i) + reserve_mask[static_cast(i)] = 1; + } + const int32_t prefix_limit = + prompt_end_position_ > 0 ? prompt_end_position_ : planned_prompt_length_; + if ((config_.protect_prefill || !config_.count_prompt_tokens) && prefix_limit > 0) { + for (int32_t i = 0; i < total_tokens; ++i) { + if (cache_positions_[static_cast(i)] < prefix_limit) + reserve_mask[static_cast(i)] = 1; + } + } + return reserve_mask; +} + +std::vector +Smollm3TriAttentionKvCache::broadcast_indices_per_head(std::vector rows, + int32_t row_count) const { + std::sort(rows.begin(), rows.end()); + std::vector keep(static_cast(cache_head_count_ * row_count)); + for (int32_t cache_head = 0; cache_head < cache_head_count_; ++cache_head) { + std::copy(rows.begin(), rows.begin() + row_count, + keep.begin() + static_cast(cache_head * row_count)); + } + return keep; +} + +std::vector +Smollm3TriAttentionKvCache::select_keep_indices(int32_t keep_budget, + Smollm3TriAttentionCompactionProfile* profile) { + const int32_t total_tokens = static_cast(cache_positions_.size()); + const int32_t old_budget = std::min(std::max(keep_budget, 0), total_tokens); + if (total_tokens <= old_budget) { + std::vector identity(static_cast(total_tokens)); + std::iota(identity.begin(), identity.end(), 0); + return broadcast_indices_per_head(std::move(identity), total_tokens); + } + + const auto reserve_mask = build_reserve_mask(total_tokens, old_budget); + std::vector reserved, candidates; + reserved.reserve(static_cast(total_tokens)); + candidates.reserve(static_cast(total_tokens)); + for (int32_t i = 0; i < total_tokens; ++i) { + if (reserve_mask[static_cast(i)] != 0) + reserved.push_back(i); + else + candidates.push_back(i); + } + if (profile != nullptr) { + profile->reserved_count = static_cast(reserved.size()); + profile->candidate_count = static_cast(candidates.size()); + } + + if (static_cast(reserved.size()) >= old_budget) { + reserved.resize(static_cast(old_budget)); + return broadcast_indices_per_head(std::move(reserved), old_budget); + } + if (candidates.empty()) + return broadcast_indices_per_head(std::move(reserved), old_budget); + +#ifdef TRTMC_HAS_CUDA_KERNELS + if (can_use_gpu_selection()) + return select_keep_indices_gpu(old_budget, reserved, candidates, profile); +#endif + return select_keep_indices_host(old_budget, reserved, candidates, profile); +} + +int32_t Smollm3TriAttentionKvCache::compaction_keep_budget(int32_t total_tokens) const { + const int32_t logical_budget = std::min(std::max(config_.kv_budget, 0), max_length_); + int32_t total_budget = logical_budget; + if (!config_.count_prompt_tokens) + total_budget = std::min(max_length_, logical_budget + count_prefix_rows()); + if (!config_.count_prompt_tokens || prompt_end_position_ > 0 || + planned_prompt_length_ <= absolute_position_) { + return std::min(total_budget, total_tokens); + } + + const int32_t remaining_prompt_tokens = + std::max(planned_prompt_length_ - absolute_position_, 0); + const int32_t slack_budget = std::clamp(max_length_ - remaining_prompt_tokens, 0, max_length_); + return std::min(std::max(total_budget, slack_budget), total_tokens); +} + +int32_t Smollm3TriAttentionKvCache::compaction_trigger_length() const { + int32_t base_trigger = config_.kv_budget; + int32_t slack_trigger = config_.kv_budget + std::max(config_.divide_length, 1); + if (!config_.count_prompt_tokens) { + const int32_t prefix_rows = count_prefix_rows(); + base_trigger += prefix_rows; + slack_trigger += prefix_rows; + } + return std::min(max_length_, std::max(base_trigger, slack_trigger)); +} + +std::vector Smollm3TriAttentionKvCache::broadcast_reserved_for_empty_budget( + int32_t keep_budget, const std::vector& reserved) const { + std::vector keep(static_cast(cache_head_count_ * keep_budget)); + for (int32_t cache_head = 0; cache_head < cache_head_count_; ++cache_head) + std::copy(reserved.begin(), reserved.end(), + keep.begin() + static_cast(cache_head * keep_budget)); + return keep; +} + +void Smollm3TriAttentionKvCache::precompute_trig_phases( + std::vector>& cos_phase, std::vector>& sin_phase, + int32_t half_dim, Smollm3TriAttentionCompactionProfile* profile) const { + if (config_.disable_trig) + return; + const auto trig_start = Clock::now(); + cos_phase.assign(offsets_.size(), std::vector(static_cast(half_dim))); + sin_phase.assign(offsets_.size(), std::vector(static_cast(half_dim))); + // TriAttention scoring must use the true absolute decode position, not the + // compacted cache length. Once earlier compactions have dropped rows, + // total_tokens no longer matches the model's current RoPE position, and + // reusing it here corrupts later-round scoring. + const float round_start = static_cast(absolute_position_); + for (std::size_t o = 0; o < offsets_.size(); ++o) { + for (int32_t d = 0; d < half_dim; ++d) { + const float phase = + (round_start + offsets_[o]) * stats_.inv_freq[static_cast(d)]; + cos_phase[o][static_cast(d)] = std::cos(phase); + sin_phase[o][static_cast(d)] = std::sin(phase); + } + } + if (profile != nullptr) + profile->trig_prep_ms += elapsed_ms(trig_start); +} + +bool Smollm3TriAttentionKvCache::layer_stats_shapes_valid( + const Smollm3TriAttentionHeadStats& layer_stats, int32_t half_dim) const { + if (layer_stats.q_mean_real.empty() || layer_stats.q_mean_imag.empty() || + layer_stats.q_abs_mean.empty() || layer_stats.freq_scale_sq.empty()) + return false; + const auto expected_stats_size = + static_cast(stats_.stats_head_count) * static_cast(half_dim); + return layer_stats.q_mean_real.size() == expected_stats_size && + layer_stats.q_mean_imag.size() == expected_stats_size && + layer_stats.q_abs_mean.size() == expected_stats_size && + layer_stats.freq_scale_sq.size() == expected_stats_size; +} + +void Smollm3TriAttentionKvCache::extract_k_rot(const float* row_ptr, int32_t head_offset, int32_t d, + int32_t half_dim, float& k_rot_real, + float& k_rot_imag) const { + if (stats_.rope_style == Smollm3TriAttentionRopeStyle::kInterleaved) { + k_rot_real = row_ptr[head_offset + (2 * d)]; + k_rot_imag = row_ptr[head_offset + (2 * d) + 1]; + } else { + k_rot_real = row_ptr[head_offset + d]; + k_rot_imag = row_ptr[head_offset + half_dim + d]; + } +} + +float Smollm3TriAttentionKvCache::reduce_trig_sums(const std::vector& trig_sums) const { + if (trig_sums.empty()) + return 0.0F; + if (config_.score_aggregation == Smollm3TriAttentionScoreAggregation::kMax) + return *std::max_element(trig_sums.begin(), trig_sums.end()); + const float sum = std::accumulate(trig_sums.begin(), trig_sums.end(), 0.0F); + return sum / static_cast(trig_sums.size()); +} + +float Smollm3TriAttentionKvCache::score_one_row( + const float* row_ptr, const Smollm3TriAttentionHeadStats& layer_stats, std::size_t stats_base, + int32_t head_offset, int32_t half_dim, const std::vector>& cos_phase, + const std::vector>& sin_phase) const { + float additive = 0.0F; + std::vector trig_sums(config_.disable_trig ? 0U : offsets_.size(), 0.0F); + for (int32_t d = 0; d < half_dim; ++d) { + float k_rot_real = 0.0F; + float k_rot_imag = 0.0F; + extract_k_rot(row_ptr, head_offset, d, half_dim, k_rot_real, k_rot_imag); + const auto idx = stats_base + static_cast(d); + const float q_real = layer_stats.q_mean_real[idx]; + const float q_imag = layer_stats.q_mean_imag[idx]; + const float q_abs = layer_stats.q_abs_mean[idx]; + const float freq_scale_sq = layer_stats.freq_scale_sq[idx]; + const float q_mean_abs = complex_abs(q_real, q_imag); + const float k_abs = complex_abs(k_rot_real, k_rot_imag); + const float extra_coef = config_.disable_mlr ? q_abs : (q_abs - q_mean_abs); + additive += k_abs * extra_coef * freq_scale_sq; + if (config_.disable_trig) + continue; + const float prod_real = q_real * k_rot_real + q_imag * k_rot_imag; + const float prod_imag = q_imag * k_rot_real - q_real * k_rot_imag; + for (std::size_t o = 0; o < offsets_.size(); ++o) { + trig_sums[o] += freq_scale_sq * (prod_real * cos_phase[o][static_cast(d)] - + prod_imag * sin_phase[o][static_cast(d)]); + } + } + return reduce_trig_sums(trig_sums) + additive; +} + +void Smollm3TriAttentionKvCache::score_rows_for_head( + std::vector& scores, const std::vector& layer_cache, + const Smollm3TriAttentionHeadStats& layer_stats, int32_t score_head, int32_t cache_head, + int32_t half_dim, int32_t total_tokens, const std::vector>& cos_phase, + const std::vector>& sin_phase) const { + const int32_t head_offset = cache_head * query_group_size_ * stats_.head_dim; + const auto stats_base = + static_cast(score_head) * static_cast(half_dim); + for (int32_t row = 0; row < total_tokens; ++row) { + const float* row_ptr = + layer_cache.data() + static_cast(row) * static_cast(kv_dim_); + scores[static_cast(row)] = score_one_row( + row_ptr, layer_stats, stats_base, head_offset, half_dim, cos_phase, sin_phase); + } +} + +void Smollm3TriAttentionKvCache::standardize_scores(std::vector& scores) const { + if (scores.empty()) + return; + const float mean = + std::accumulate(scores.begin(), scores.end(), 0.0F) / static_cast(scores.size()); + const float stddev = score_full_stddev_or_one(scores, mean); + for (float& value : scores) + value = (value - mean) / stddev; +} + +void Smollm3TriAttentionKvCache::accumulate_layer_fallback( + const std::vector>& layer_scores, std::vector& global_fallback_sum, + int32_t& global_fallback_count, int32_t total_tokens) const { + for (const auto& scores : layer_scores) { + for (int32_t row = 0; row < total_tokens; ++row) + global_fallback_sum[static_cast(row)] += + scores[static_cast(row)]; + ++global_fallback_count; + } +} + +void Smollm3TriAttentionKvCache::reduce_group_into_aggregate( + int32_t cache_head, int32_t total_tokens, const std::vector>& layer_scores, + const std::vector& sampled_group, std::vector& aggregate, + bool first_layer_for_cache_head, float* layer_dump) const { + const bool use_max = config_.per_layer_aggregation == Smollm3TriAttentionScoreAggregation::kMax; + (void)cache_head; + for (int32_t row = 0; row < total_tokens; ++row) { + float reduced = layer_scores[static_cast(sampled_group.front())] + [static_cast(row)]; + for (std::size_t group_idx = 1; group_idx < sampled_group.size(); ++group_idx) { + reduced = + std::max(reduced, layer_scores[static_cast(sampled_group[group_idx])] + [static_cast(row)]); + } + if (use_max) { + aggregate[static_cast(row)] = + first_layer_for_cache_head + ? reduced + : std::max(aggregate[static_cast(row)], reduced); + } else { + aggregate[static_cast(row)] += reduced; + } + if (layer_dump != nullptr) + layer_dump[static_cast(row)] = reduced; + } +} + +void Smollm3TriAttentionKvCache::accumulate_layer_to_aggregate( + int32_t layer, int32_t total_tokens, const std::vector>& layer_scores, + const std::vector>& sampled_by_cache_head, + std::vector>& aggregated_scores, + std::vector& contributing_layers_by_cache_head, + std::vector& layer_aggregate_dump) const { + for (int32_t cache_head = 0; cache_head < cache_head_count_; ++cache_head) { + const auto& sampled_group = sampled_by_cache_head[static_cast(cache_head)]; + if (sampled_group.empty()) + continue; + const bool first_layer_for_cache_head = + contributing_layers_by_cache_head[static_cast(cache_head)] == 0; + float* layer_dump = nullptr; + if (!layer_aggregate_dump.empty()) { + layer_dump = + layer_aggregate_dump.data() + + (static_cast(layer) * static_cast(cache_head_count_) + + static_cast(cache_head)) * + static_cast(total_tokens); + } + reduce_group_into_aggregate(cache_head, total_tokens, layer_scores, sampled_group, + aggregated_scores[static_cast(cache_head)], + first_layer_for_cache_head, layer_dump); + ++contributing_layers_by_cache_head[static_cast(cache_head)]; + } +} + +std::vector +Smollm3TriAttentionKvCache::compute_fallback_mean(const std::vector& global_fallback_sum, + int32_t global_fallback_count, + int32_t total_tokens) const { + std::vector mean(static_cast(total_tokens), 0.0F); + if (global_fallback_count <= 0) + return mean; + const float inv_count = 1.0F / static_cast(global_fallback_count); + for (int32_t row = 0; row < total_tokens; ++row) + mean[static_cast(row)] = + global_fallback_sum[static_cast(row)] * inv_count; + return mean; +} + +void Smollm3TriAttentionKvCache::finalize_per_head_aggregate( + std::vector>& aggregated_scores, + const std::vector& contributing_layers_by_cache_head, + const std::vector& global_fallback_mean) const { + const bool per_layer_max = + config_.per_layer_aggregation == Smollm3TriAttentionScoreAggregation::kMax; + for (int32_t cache_head = 0; cache_head < cache_head_count_; ++cache_head) { + auto& scores = aggregated_scores[static_cast(cache_head)]; + const int32_t head_layer_count = + contributing_layers_by_cache_head[static_cast(cache_head)]; + if (head_layer_count <= 0) { + scores = global_fallback_mean; + continue; + } + if (per_layer_max) + continue; + const float inv = 1.0F / static_cast(head_layer_count); + for (float& value : scores) + value *= inv; + } +} + +void Smollm3TriAttentionKvCache::maybe_dump_score_values( + const std::vector>& aggregated_scores, + const std::vector& layer_aggregate_dump, int32_t total_tokens) const { + if (!config_.dump_score_values) + return; + const char* dump_path = + config_.dump_keep_path.empty() ? nullptr : config_.dump_keep_path.c_str(); + if (dump_path == nullptr || dump_path[0] == '\0') + return; + std::vector packed_aggregate( + static_cast(cache_head_count_) * static_cast(total_tokens), 0.0F); + for (int32_t cache_head = 0; cache_head < cache_head_count_; ++cache_head) { + std::copy(aggregated_scores[static_cast(cache_head)].begin(), + aggregated_scores[static_cast(cache_head)].end(), + packed_aggregate.begin() + static_cast(cache_head) * + static_cast(total_tokens)); + } + char path_buf[1024]; + std::snprintf(path_buf, sizeof(path_buf), "%s.agg.bin", dump_path); + std::ofstream out_agg(path_buf, std::ios::binary); + out_agg.write(reinterpret_cast(packed_aggregate.data()), + static_cast(packed_aggregate.size() * sizeof(float))); + if (layer_aggregate_dump.empty()) + return; + std::snprintf(path_buf, sizeof(path_buf), "%s.layeragg.bin", dump_path); + std::ofstream out_layer(path_buf, std::ios::binary); + out_layer.write(reinterpret_cast(layer_aggregate_dump.data()), + static_cast(layer_aggregate_dump.size() * sizeof(float))); +} + +std::vector Smollm3TriAttentionKvCache::build_keep_indices_per_head( + const std::vector>& aggregated_scores, const std::vector& reserved, + const std::vector& candidates, int32_t keep_budget, int32_t need) const { + std::vector keep(static_cast(cache_head_count_ * keep_budget)); + for (int32_t cache_head = 0; cache_head < cache_head_count_; ++cache_head) { + auto* out = keep.data() + + static_cast(cache_head) * static_cast(keep_budget); + std::copy(reserved.begin(), reserved.end(), out); + std::vector> ranked; + ranked.reserve(candidates.size()); + const auto& scores = aggregated_scores[static_cast(cache_head)]; + for (const int32_t row : candidates) + ranked.emplace_back(scores[static_cast(row)], row); + const int32_t top_n = std::min(need, ranked.size()); + std::partial_sort(ranked.begin(), ranked.begin() + top_n, ranked.end(), + [](const auto& a, const auto& b) { + if (a.first != b.first) + return a.first > b.first; + return a.second < b.second; + }); + for (int32_t i = 0; i < top_n; ++i) + out[static_cast(reserved.size() + i)] = + ranked[static_cast(i)].second; + std::sort(out, out + keep_budget); + } + return keep; +} + +bool Smollm3TriAttentionKvCache::process_layer_for_host_selection( + int32_t layer, int32_t half_dim, int32_t total_tokens, + const std::vector>& cos_phase, + const std::vector>& sin_phase, + std::vector>& aggregated_scores, std::vector& global_fallback_sum, + int32_t& global_fallback_count, std::vector& contributing_layers_by_cache_head, + std::vector& layer_aggregate_dump, Smollm3TriAttentionCompactionProfile* profile) const { + if (layer < 0 || layer >= static_cast(stats_.layer_stats.size())) + return false; + const auto& layer_stats = stats_.layer_stats[static_cast(layer)]; + const auto& sampled_heads = + stats_.sampled_score_heads_by_layer[static_cast(layer)]; + if (sampled_heads.empty() || !layer_stats_shapes_valid(layer_stats, half_dim)) + return false; + + const auto layer_cache = + copy_cache_rows_to_host(cache_k_[static_cast(layer)], cache_length_, profile); + if (profile != nullptr) { + profile->sampled_layers += 1; + profile->sampled_heads += static_cast(sampled_heads.size()); + } + + std::vector> layer_scores( + sampled_heads.size(), std::vector(static_cast(total_tokens), 0.0F)); + std::vector> sampled_by_cache_head( + static_cast(cache_head_count_)); + for (int32_t sampled_idx = 0; sampled_idx < static_cast(sampled_heads.size()); + ++sampled_idx) { + const int32_t score_head = sampled_heads[static_cast(sampled_idx)]; + const int32_t cache_head = std::min(cache_head_count_ - 1, score_head / score_group_size_); + sampled_by_cache_head[static_cast(cache_head)].push_back(sampled_idx); + auto& scores = layer_scores[static_cast(sampled_idx)]; + score_rows_for_head(scores, layer_cache, layer_stats, score_head, cache_head, half_dim, + total_tokens, cos_phase, sin_phase); + standardize_scores(scores); + } + + accumulate_layer_fallback(layer_scores, global_fallback_sum, global_fallback_count, + total_tokens); + accumulate_layer_to_aggregate(layer, total_tokens, layer_scores, sampled_by_cache_head, + aggregated_scores, contributing_layers_by_cache_head, + layer_aggregate_dump); + return true; +} + +std::vector Smollm3TriAttentionKvCache::select_keep_indices_host( + int32_t keep_budget, const std::vector& reserved, + const std::vector& candidates, Smollm3TriAttentionCompactionProfile* profile) const { + const int32_t need = std::max(0, keep_budget - static_cast(reserved.size())); + if (need <= 0) + return broadcast_reserved_for_empty_budget(keep_budget, reserved); + const int32_t half_dim = stats_.head_dim / 2; + const int32_t total_tokens = static_cast(cache_positions_.size()); + if (half_dim <= 0 || total_tokens <= 0) + return {}; + std::vector> cos_phase; + std::vector> sin_phase; + precompute_trig_phases(cos_phase, sin_phase, half_dim, profile); + + const auto score_start = Clock::now(); + std::vector> aggregated_scores( + static_cast(cache_head_count_), + std::vector(static_cast(total_tokens), 0.0F)); + std::vector layer_aggregate_dump; + if (config_.dump_score_values) { + layer_aggregate_dump.assign(static_cast(num_layers_) * + static_cast(cache_head_count_) * + static_cast(total_tokens), + std::numeric_limits::quiet_NaN()); + } + std::vector global_fallback_sum(static_cast(total_tokens), 0.0F); + int32_t global_fallback_count = 0; + std::vector contributing_layers_by_cache_head( + static_cast(cache_head_count_), 0); + int32_t contributing_layers = 0; + for (int32_t layer = 0; layer < num_layers_; ++layer) { + if (process_layer_for_host_selection( + layer, half_dim, total_tokens, cos_phase, sin_phase, aggregated_scores, + global_fallback_sum, global_fallback_count, contributing_layers_by_cache_head, + layer_aggregate_dump, profile)) + ++contributing_layers; + } + if (contributing_layers <= 0) + return {}; + const auto global_fallback_mean = + compute_fallback_mean(global_fallback_sum, global_fallback_count, total_tokens); + finalize_per_head_aggregate(aggregated_scores, contributing_layers_by_cache_head, + global_fallback_mean); + if (profile != nullptr) + profile->score_ms += elapsed_ms(score_start); + maybe_dump_score_values(aggregated_scores, layer_aggregate_dump, total_tokens); + + const auto combine_start = Clock::now(); + auto keep = + build_keep_indices_per_head(aggregated_scores, reserved, candidates, keep_budget, need); + if (profile != nullptr) + profile->combine_ms += elapsed_ms(combine_start); + return keep; +} + +#ifdef TRTMC_HAS_CUDA_KERNELS +void Smollm3TriAttentionKvCache::standardize_score_rows(float* rows, int32_t num_rows, + int32_t total_tokens) const { + for (int32_t r = 0; r < num_rows; ++r) { + float* score_row = rows + static_cast(r) * total_tokens; + float mean = 0.0F; + for (int32_t row = 0; row < total_tokens; ++row) + mean += score_row[row]; + mean /= static_cast(total_tokens); + float var = 0.0F; + for (int32_t row = 0; row < total_tokens; ++row) { + const float delta = score_row[row] - mean; + var += delta * delta; + } + const float denom = total_tokens > 1 ? static_cast(total_tokens - 1) : 1.0F; + const float stddev = std::sqrt(std::max(var / denom, 0.0F)); + const float std_safe = stddev < kEps ? 1.0F : stddev; + for (int32_t row = 0; row < total_tokens; ++row) + score_row[row] = (score_row[row] - mean) / std_safe; + } +} + +void Smollm3TriAttentionKvCache::accumulate_flat_fallback(const float* rows, int32_t num_rows, + int32_t total_tokens, + std::vector& fallback_sum, + int32_t& fallback_count) const { + for (int32_t r = 0; r < num_rows; ++r) { + const float* score_row = rows + static_cast(r) * total_tokens; + for (int32_t row = 0; row < total_tokens; ++row) + fallback_sum[static_cast(row)] += score_row[row]; + ++fallback_count; + } +} + +void Smollm3TriAttentionKvCache::aggregate_gpu_layer_into_cache_heads( + const std::vector& host_scores, const LayerGpuStats& gpu, int32_t total_tokens, + std::vector& aggregated_scores, + std::vector& contributing_layers_by_cache_head) const { + const bool use_max = config_.per_layer_aggregation == Smollm3TriAttentionScoreAggregation::kMax; + for (int32_t cache_head = 0; cache_head < cache_head_count_; ++cache_head) { + std::vector sampled_group; + sampled_group.reserve(static_cast(gpu.score_head_count)); + for (int32_t score_head = 0; score_head < gpu.score_head_count; ++score_head) { + if (gpu.host_cache_head_indices[static_cast(score_head)] == cache_head) + sampled_group.push_back(score_head); + } + if (sampled_group.empty()) + continue; + float* aggregate_row = + aggregated_scores.data() + static_cast(cache_head) * total_tokens; + const bool first_layer = + contributing_layers_by_cache_head[static_cast(cache_head)] == 0; + for (int32_t row = 0; row < total_tokens; ++row) { + float reduced = + host_scores[static_cast(sampled_group.front()) * total_tokens + + static_cast(row)]; + for (std::size_t group_idx = 1; group_idx < sampled_group.size(); ++group_idx) { + reduced = std::max( + reduced, + host_scores[static_cast(sampled_group[group_idx]) * total_tokens + + static_cast(row)]); + } + aggregate_row[row] = + use_max ? (first_layer ? reduced : std::max(aggregate_row[row], reduced)) + : aggregate_row[row] + reduced; + } + ++contributing_layers_by_cache_head[static_cast(cache_head)]; + } +} + +Smollm3TriAttentionKvCache::GpuLayerResult +Smollm3TriAttentionKvCache::process_layer_for_gpu_selection( + int32_t layer, int32_t total_tokens, int32_t num_offsets, std::vector& aggregated_scores, + std::vector& global_fallback_sum, int32_t& global_fallback_count, + std::vector& contributing_layers_by_cache_head, + Smollm3TriAttentionCompactionProfile* profile) { + if (layer < 0 || layer >= static_cast(layer_gpu_stats_.size())) + return GpuLayerResult::kSkipped; + auto& gpu = layer_gpu_stats_[static_cast(layer)]; + if (gpu.score_head_count <= 0) + return GpuLayerResult::kSkipped; + if (profile != nullptr) { + profile->sampled_layers += 1; + profile->sampled_heads += gpu.score_head_count; + } + const bool launched = smollm3_triattention_score_candidates_gpu( + cache_k_[static_cast(layer)].data(), cache_dtype_, kv_dim_, stats_.head_dim, + (stats_.rope_style == Smollm3TriAttentionRopeStyle::kInterleaved), + static_cast(candidate_indices_device_.data()), total_tokens, nullptr, + static_cast(inv_freq_device_.data()), + static_cast(cos_phase_device_.data()), + static_cast(sin_phase_device_.data()), num_offsets, + static_cast(gpu.head_offsets.data()), + static_cast(gpu.head_cache_indices.data()), + static_cast(gpu.q_mean_real.data()), + static_cast(gpu.q_mean_imag.data()), + static_cast(gpu.q_abs_mean.data()), + static_cast(gpu.freq_scale_sq.data()), gpu.score_head_count, + config_.disable_mlr, config_.disable_trig, + config_.score_aggregation == Smollm3TriAttentionScoreAggregation::kMax, + static_cast(gpu.scores.data()), stream_); + if (!launched) + return GpuLayerResult::kFailed; + std::vector host_scores(static_cast(gpu.score_head_count) * + static_cast(total_tokens)); + const auto score_bytes = host_scores.size() * sizeof(float); + if (cudaMemcpyAsync(host_scores.data(), gpu.scores.data(), score_bytes, cudaMemcpyDeviceToHost, + stream_) != cudaSuccess) + return GpuLayerResult::kFailed; + if (cudaStreamSynchronize(stream_) != cudaSuccess) + return GpuLayerResult::kFailed; + standardize_score_rows(host_scores.data(), gpu.score_head_count, total_tokens); + accumulate_flat_fallback(host_scores.data(), gpu.score_head_count, total_tokens, + global_fallback_sum, global_fallback_count); + aggregate_gpu_layer_into_cache_heads(host_scores, gpu, total_tokens, aggregated_scores, + contributing_layers_by_cache_head); + return GpuLayerResult::kContributed; +} + +bool Smollm3TriAttentionKvCache::upload_candidate_indices_identity(int32_t total_tokens) { + std::vector score_rows(static_cast(total_tokens)); + std::iota(score_rows.begin(), score_rows.end(), 0); + const auto bytes = static_cast(total_tokens) * sizeof(int32_t); + return cudaMemcpyAsync(candidate_indices_device_.data(), score_rows.data(), bytes, + cudaMemcpyHostToDevice, stream_) == cudaSuccess; +} + +bool Smollm3TriAttentionKvCache::upload_gpu_trig_phases( + int32_t num_offsets, int32_t half_dim, Smollm3TriAttentionCompactionProfile* profile) { + if (config_.disable_trig) + return true; + const auto trig_start = Clock::now(); + std::vector cos_phase(static_cast(num_offsets) * + static_cast(half_dim)); + std::vector sin_phase(static_cast(num_offsets) * + static_cast(half_dim)); + const float round_start = static_cast(absolute_position_); + for (int32_t o = 0; o < num_offsets; ++o) { + for (int32_t d = 0; d < half_dim; ++d) { + const std::size_t idx = + static_cast(o) * static_cast(half_dim) + d; + const float phase = (round_start + offsets_[static_cast(o)]) * + stats_.inv_freq[static_cast(d)]; + cos_phase[idx] = std::cos(phase); + sin_phase[idx] = std::sin(phase); + } + } + const auto phase_bytes = cos_phase.size() * sizeof(float); + const bool ok_cos = cudaMemcpyAsync(cos_phase_device_.data(), cos_phase.data(), phase_bytes, + cudaMemcpyHostToDevice, stream_) == cudaSuccess; + const bool ok_sin = cudaMemcpyAsync(sin_phase_device_.data(), sin_phase.data(), phase_bytes, + cudaMemcpyHostToDevice, stream_) == cudaSuccess; + if (profile != nullptr) + profile->trig_prep_ms += elapsed_ms(trig_start); + return ok_cos && ok_sin; +} + +void Smollm3TriAttentionKvCache::finalize_flat_per_head_aggregate( + std::vector& aggregated_scores, + const std::vector& contributing_layers_by_cache_head, + const std::vector& global_fallback_mean, int32_t total_tokens) const { + const bool per_layer_max = + config_.per_layer_aggregation == Smollm3TriAttentionScoreAggregation::kMax; + for (int32_t cache_head = 0; cache_head < cache_head_count_; ++cache_head) { + float* score_row = + aggregated_scores.data() + static_cast(cache_head) * total_tokens; + const int32_t head_layer_count = + contributing_layers_by_cache_head[static_cast(cache_head)]; + if (head_layer_count <= 0) { + std::copy(global_fallback_mean.begin(), global_fallback_mean.end(), score_row); + continue; + } + if (per_layer_max) + continue; + const float inv = 1.0F / static_cast(head_layer_count); + for (int32_t row = 0; row < total_tokens; ++row) + score_row[row] *= inv; + } +} + +std::vector Smollm3TriAttentionKvCache::build_keep_from_flat_aggregate( + const std::vector& aggregated_scores, const std::vector& reserved, + const std::vector& candidates, int32_t keep_budget, int32_t need, + int32_t total_tokens) const { + std::vector keep(static_cast(cache_head_count_ * keep_budget)); + for (int32_t cache_head = 0; cache_head < cache_head_count_; ++cache_head) { + const float* score_row = + aggregated_scores.data() + static_cast(cache_head) * total_tokens; + auto* out = keep.data() + + static_cast(cache_head) * static_cast(keep_budget); + std::copy(reserved.begin(), reserved.end(), out); + std::vector> ranked; + ranked.reserve(candidates.size()); + for (const int32_t row : candidates) + ranked.emplace_back(score_row[row], row); + const int32_t top_n = std::min(need, ranked.size()); + std::partial_sort(ranked.begin(), ranked.begin() + top_n, ranked.end(), + [](const auto& a, const auto& b) { + if (a.first != b.first) + return a.first > b.first; + return a.second < b.second; + }); + for (int32_t i = 0; i < top_n; ++i) + out[static_cast(reserved.size() + i)] = + ranked[static_cast(i)].second; + std::sort(out, out + keep_budget); + } + return keep; +} + +bool Smollm3TriAttentionKvCache::run_gpu_selection_over_layers( + int32_t total_tokens, int32_t num_offsets, std::vector& aggregated_scores, + std::vector& global_fallback_sum, int32_t& global_fallback_count, + std::vector& contributing_layers_by_cache_head, int32_t& contributing_layers, + Smollm3TriAttentionCompactionProfile* profile) { + for (int32_t layer = 0; layer < num_layers_; ++layer) { + const auto status = process_layer_for_gpu_selection( + layer, total_tokens, num_offsets, aggregated_scores, global_fallback_sum, + global_fallback_count, contributing_layers_by_cache_head, profile); + if (status == GpuLayerResult::kFailed) + return false; + if (status == GpuLayerResult::kContributed) + ++contributing_layers; + } + return true; +} + +std::vector Smollm3TriAttentionKvCache::select_keep_indices_gpu( + int32_t keep_budget, const std::vector& reserved, + const std::vector& candidates, Smollm3TriAttentionCompactionProfile* profile) { + const int32_t need = std::max(0, keep_budget - static_cast(reserved.size())); + if (need <= 0) + return broadcast_reserved_for_empty_budget(keep_budget, reserved); + if (candidates.empty()) + return select_keep_indices_host(keep_budget, reserved, candidates, profile); + + const int32_t total_tokens = static_cast(cache_positions_.size()); + if (total_tokens <= 0) + return {}; + if (!upload_candidate_indices_identity(total_tokens)) + return select_keep_indices_host(keep_budget, reserved, candidates, profile); + + const int32_t num_offsets = static_cast(offsets_.size()); + const int32_t half_dim = stats_.head_dim / 2; + if (!upload_gpu_trig_phases(num_offsets, half_dim, profile)) + return select_keep_indices_host(keep_budget, reserved, candidates, profile); + + const auto score_start = Clock::now(); + std::vector aggregated_scores( + static_cast(cache_head_count_) * static_cast(total_tokens), 0.0F); + std::vector global_fallback_sum(static_cast(total_tokens), 0.0F); + int32_t global_fallback_count = 0; + std::vector contributing_layers_by_cache_head( + static_cast(cache_head_count_), 0); + int32_t contributing_layers = 0; + if (!run_gpu_selection_over_layers( + total_tokens, num_offsets, aggregated_scores, global_fallback_sum, + global_fallback_count, contributing_layers_by_cache_head, contributing_layers, profile)) + return select_keep_indices_host(keep_budget, reserved, candidates, profile); + if (contributing_layers <= 0) + return select_keep_indices_host(keep_budget, reserved, candidates, profile); + + const auto global_fallback_mean = + compute_fallback_mean(global_fallback_sum, global_fallback_count, total_tokens); + finalize_flat_per_head_aggregate(aggregated_scores, contributing_layers_by_cache_head, + global_fallback_mean, total_tokens); + if (profile != nullptr) + profile->score_ms += elapsed_ms(score_start); + + const auto combine_start = Clock::now(); + auto keep = build_keep_from_flat_aggregate(aggregated_scores, reserved, candidates, keep_budget, + need, total_tokens); + if (profile != nullptr) + profile->combine_ms += elapsed_ms(combine_start); + return keep; +} +#endif + +void Smollm3TriAttentionKvCache::dump_cache_rows_to_file(const DeviceTensor& tensor, int32_t rows, + const std::string& path, + std::vector& out_files) { + const auto host_cache = copy_cache_rows_to_host(tensor, rows, nullptr); + std::vector packed(static_cast(rows) * + static_cast(cache_head_count_) * + static_cast(stats_.head_dim), + 0.0F); + for (int32_t row = 0; row < rows; ++row) { + const float* row_ptr = + host_cache.data() + static_cast(row) * static_cast(kv_dim_); + for (int32_t cache_head = 0; cache_head < cache_head_count_; ++cache_head) { + const int32_t head_offset = cache_head * query_group_size_ * stats_.head_dim; + float* dst = packed.data() + (static_cast(row) * + static_cast(cache_head_count_) + + static_cast(cache_head)) * + static_cast(stats_.head_dim); + std::copy_n(row_ptr + head_offset, stats_.head_dim, dst); + } + } + std::ofstream score_out(path, std::ios::binary); + score_out.write(reinterpret_cast(packed.data()), + static_cast(packed.size() * sizeof(float))); + out_files.emplace_back(path); +} + +void Smollm3TriAttentionKvCache::maybe_dump_score_cache(int32_t rows, const char* k_pattern, + const char* v_pattern, + std::vector& score_files, + std::vector& value_files) { + if (!config_.dump_score_cache) + return; + const char* dump_path = + config_.dump_keep_path.empty() ? nullptr : config_.dump_keep_path.c_str(); + if (dump_path == nullptr || dump_path[0] == '\0') + return; + char path_buf[1024]; + for (int32_t layer = 0; layer < num_layers_; ++layer) { + std::snprintf(path_buf, sizeof(path_buf), k_pattern, dump_path, layer); + dump_cache_rows_to_file(cache_k_[static_cast(layer)], rows, path_buf, + score_files); + std::snprintf(path_buf, sizeof(path_buf), v_pattern, dump_path, layer); + dump_cache_rows_to_file(cache_v_[static_cast(layer)], rows, path_buf, + value_files); + } +} + +#ifdef TRTMC_HAS_CUDA_KERNELS +bool Smollm3TriAttentionKvCache::gpu_compaction_upload_keep( + int32_t keep_count, const std::vector& keep_indices) { + if (config_.disable_gpu_compaction || keep_count <= 0 || !keep_indices_device_.ok()) + return false; + if (!scratch_k_device_.ok() || !scratch_v_device_.ok()) + return false; + const auto keep_bytes = static_cast(keep_indices.size()) * sizeof(int32_t); + return cudaMemcpyAsync(keep_indices_device_.data(), keep_indices.data(), keep_bytes, + cudaMemcpyHostToDevice, stream_) == cudaSuccess; +} + +bool Smollm3TriAttentionKvCache::gpu_compact_one_layer(int32_t layer, int32_t keep_count, + std::size_t row_bytes) { + const bool ok_k = smollm3_triattention_compact_rows_gpu( + cache_k_[static_cast(layer)].data(), scratch_k_device_.data(), cache_dtype_, + kv_dim_, static_cast(keep_indices_device_.data()), keep_count, + stats_.head_dim, cache_head_count_, query_group_size_, stream_); + const bool ok_v = smollm3_triattention_compact_rows_gpu( + cache_v_[static_cast(layer)].data(), scratch_v_device_.data(), cache_dtype_, + kv_dim_, static_cast(keep_indices_device_.data()), keep_count, + stats_.head_dim, cache_head_count_, query_group_size_, stream_); + if (!ok_k || !ok_v) + return false; + const auto bytes = static_cast(keep_count) * row_bytes; + const bool copy_k = + cudaMemcpyAsync(cache_k_[static_cast(layer)].data(), scratch_k_device_.data(), + bytes, cudaMemcpyDeviceToDevice, stream_) == cudaSuccess; + const bool copy_v = + cudaMemcpyAsync(cache_v_[static_cast(layer)].data(), scratch_v_device_.data(), + bytes, cudaMemcpyDeviceToDevice, stream_) == cudaSuccess; + return copy_k && copy_v; +} +#endif + +bool Smollm3TriAttentionKvCache::compact_layer_on_gpu(int32_t layer, + const std::vector& keep_indices, + int32_t keep_count, std::size_t row_bytes, + int64_t& repack_calls, + std::size_t& repack_bytes) { +#ifdef TRTMC_HAS_CUDA_KERNELS + if (!gpu_compaction_upload_keep(keep_count, keep_indices)) + return false; + if (!gpu_compact_one_layer(layer, keep_count, row_bytes)) + return false; + repack_calls += 4; + repack_bytes += static_cast(keep_count) * row_bytes * 2U; + return true; +#else + (void)layer; + (void)keep_indices; + (void)keep_count; + (void)row_bytes; + (void)repack_calls; + (void)repack_bytes; + return false; +#endif +} + +void Smollm3TriAttentionKvCache::compact_layer_on_host( + int32_t layer, const std::vector& keep_indices, int32_t keep_count, + std::size_t row_bytes, std::size_t head_block_bytes, int32_t old_cache_length, + int64_t& repack_calls, std::size_t& repack_bytes) { + auto* ck = static_cast(cache_k_[static_cast(layer)].data()); + auto* cv = static_cast(cache_v_[static_cast(layer)].data()); + for (int32_t dst = 0; dst < keep_count; ++dst) { + const auto dst_offset = static_cast(dst) * row_bytes; + for (int32_t cache_head = 0; cache_head < cache_head_count_; ++cache_head) { + const int32_t src = + keep_indices[static_cast(cache_head * keep_count + dst)]; + const auto src_offset = static_cast(src) * row_bytes + + static_cast(cache_head) * head_block_bytes; + const auto head_offset = + dst_offset + static_cast(cache_head) * head_block_bytes; + cudaMemcpyAsync(ck + head_offset, ck + src_offset, head_block_bytes, + cudaMemcpyDeviceToDevice, stream_); + cudaMemcpyAsync(cv + head_offset, cv + src_offset, head_block_bytes, + cudaMemcpyDeviceToDevice, stream_); + repack_calls += 2; + repack_bytes += head_block_bytes * 2U; + } + } + if (config_.zero_tail && old_cache_length > keep_count) { + const auto tail_offset = static_cast(keep_count) * row_bytes; + const auto tail_bytes = static_cast(old_cache_length - keep_count) * row_bytes; + cudaMemsetAsync(ck + tail_offset, 0, tail_bytes, stream_); + cudaMemsetAsync(cv + tail_offset, 0, tail_bytes, stream_); + } +} + +void Smollm3TriAttentionKvCache::compact_existing_cache(bool reserve_slot_for_append) { + const int32_t trigger_length = compaction_trigger_length(); + if (cache_length_ < trigger_length) + return; + + Smollm3TriAttentionCompactionProfile profile; + Smollm3TriAttentionCompactionProfile* profile_ptr = profile_enabled_ ? &profile : nullptr; + const auto row_bytes = static_cast(kv_dim_) * cache_element_size_; + const auto head_block_bytes = + static_cast(query_group_size_ * stats_.head_dim) * cache_element_size_; + const int32_t old_cache_length = cache_length_; + int32_t keep_count = compaction_keep_budget(cache_length_); + if (reserve_slot_for_append) + keep_count = std::min(keep_count, std::max(0, max_length_ - 1)); + cudaEvent_t repack_start = nullptr; + cudaEvent_t repack_stop = nullptr; + if (profile_ptr != nullptr) { + cudaEventCreate(&repack_start); + cudaEventCreate(&repack_stop); + cudaEventRecord(repack_start, stream_); + } + int64_t repack_calls = 0; + std::size_t repack_bytes = 0; + const auto select_start = Clock::now(); + std::vector keep_indices = select_keep_indices(keep_count, profile_ptr); + if (profile_ptr != nullptr) + profile_ptr->select_ms += elapsed_ms(select_start); + if (static_cast(keep_indices.size()) != cache_head_count_ * keep_count) { + throw std::runtime_error("TriAttention keep index shape mismatch during compaction"); + } + std::vector score_cache_files; + std::vector value_cache_files; + std::vector post_score_cache_files; + std::vector post_value_cache_files; + maybe_dump_score_cache(static_cast(cache_positions_.size()), "%s.layer%02d.bin", + "%s.v.layer%02d.bin", score_cache_files, value_cache_files); + for (int32_t layer = 0; layer < num_layers_; ++layer) { + if (compact_layer_on_gpu(layer, keep_indices, keep_count, row_bytes, repack_calls, + repack_bytes)) + continue; + compact_layer_on_host(layer, keep_indices, keep_count, row_bytes, head_block_bytes, + old_cache_length, repack_calls, repack_bytes); + } + finalize_repack_profile(profile_ptr, repack_start, repack_stop, repack_calls, repack_bytes); + maybe_dump_score_cache(keep_count, "%s.post.layer%02d.bin", "%s.post.v.layer%02d.bin", + post_score_cache_files, post_value_cache_files); + + auto new_positions_by_head = build_new_positions_by_head(keep_indices, keep_count); + log_compact_debug(new_positions_by_head.front(), keep_indices, keep_count); + ++compaction_count_; + log_compact_profile(profile_ptr, keep_count); + emit_keep_dump_json(keep_indices, keep_count, new_positions_by_head, profile_ptr, + score_cache_files, value_cache_files, post_score_cache_files, + post_value_cache_files); + cache_positions_per_head_ = std::move(new_positions_by_head); + sync_shared_positions_from_head0(); + cache_length_ = keep_count; +} + +void Smollm3TriAttentionKvCache::finalize_repack_profile( + Smollm3TriAttentionCompactionProfile* profile_ptr, cudaEvent_t repack_start, + cudaEvent_t repack_stop, int64_t repack_calls, std::size_t repack_bytes) const { + if (profile_ptr == nullptr) + return; + cudaEventRecord(repack_stop, stream_); + cudaEventSynchronize(repack_stop); + float repack_ms = 0.0F; + cudaEventElapsedTime(&repack_ms, repack_start, repack_stop); + profile_ptr->repack_ms = static_cast(repack_ms); + profile_ptr->repack_calls = repack_calls; + profile_ptr->repack_bytes = repack_bytes; + cudaEventDestroy(repack_start); + cudaEventDestroy(repack_stop); +} + +std::vector> +Smollm3TriAttentionKvCache::build_new_positions_by_head(const std::vector& keep_indices, + int32_t keep_count) const { + std::vector> out(static_cast(cache_head_count_)); + for (int32_t cache_head = 0; cache_head < cache_head_count_; ++cache_head) { + auto& head_out = out[static_cast(cache_head)]; + const auto& old_positions = cache_positions_per_head_[static_cast(cache_head)]; + head_out.reserve(static_cast(keep_count)); + for (int32_t dst = 0; dst < keep_count; ++dst) { + const int32_t idx = + keep_indices[static_cast(cache_head * keep_count + dst)]; + head_out.push_back(old_positions[static_cast(idx)]); + } + } + return out; +} + +std::vector +Smollm3TriAttentionKvCache::collect_dropped_positions(const std::vector& keep_indices, + int32_t keep_count) const { + std::vector keep_mask(static_cast(cache_length_), 0); + for (int32_t dst = 0; dst < keep_count; ++dst) { + const int32_t idx = keep_indices[static_cast(dst)]; + keep_mask[static_cast(idx)] = 1; + } + std::vector dropped; + for (int32_t i = 0; i < cache_length_; ++i) { + if (keep_mask[static_cast(i)] == 0) + dropped.push_back(cache_positions_[static_cast(i)]); + } + return dropped; +} + +int32_t Smollm3TriAttentionKvCache::count_prefix_positions( + const std::vector& representative_positions) const { + const int32_t prefix_limit = + prompt_end_position_ > 0 ? prompt_end_position_ : planned_prompt_length_; + if (prefix_limit <= 0) + return 0; + int32_t kept = 0; + for (int32_t pos : representative_positions) { + if (pos < prefix_limit) + ++kept; + } + return kept; +} + +void Smollm3TriAttentionKvCache::log_compact_debug( + const std::vector& representative_positions, const std::vector& keep_indices, + int32_t keep_count) const { + if (!config_.debug) + return; + const int32_t kept_prefix = count_prefix_positions(representative_positions); + const auto dropped_positions = collect_dropped_positions(keep_indices, keep_count); + const int32_t first_pos = + representative_positions.empty() ? -1 : representative_positions.front(); + const int32_t last_pos = + representative_positions.empty() ? -1 : representative_positions.back(); + std::cerr << "[trtmc.triattention] compact abs_pos=" << absolute_position_ + << " old_rows=" << cache_length_ << " kept_rows=" << keep_count + << " kept_prefix=" << kept_prefix << " first_pos=" << first_pos + << " last_pos=" << last_pos; + if (!dropped_positions.empty()) { + std::cerr << " dropped_pos="; + for (std::size_t i = 0; i < dropped_positions.size(); ++i) { + if (i > 0) + std::cerr << ','; + std::cerr << dropped_positions[i]; + } + } + std::cerr << '\n'; +} + +void Smollm3TriAttentionKvCache::log_compact_profile( + const Smollm3TriAttentionCompactionProfile* profile_ptr, int32_t keep_count) const { + if (profile_ptr == nullptr) + return; + std::cerr << "[trtmc.triattention.profile] compact#" << compaction_count_ + << " abs_pos=" << absolute_position_ << " old_rows=" << cache_length_ + << " kept_rows=" << keep_count << " reserved=" << profile_ptr->reserved_count + << " candidates=" << profile_ptr->candidate_count + << " sampled_layers=" << profile_ptr->sampled_layers + << " sampled_heads=" << profile_ptr->sampled_heads + << " host_copy_ms=" << profile_ptr->host_copy_ms + << " host_convert_ms=" << profile_ptr->host_convert_ms + << " trig_prep_ms=" << profile_ptr->trig_prep_ms + << " score_ms=" << profile_ptr->score_ms << " combine_ms=" << profile_ptr->combine_ms + << " select_ms=" << profile_ptr->select_ms << " repack_ms=" << profile_ptr->repack_ms + << " host_copy_mb=" + << (static_cast(profile_ptr->host_copy_bytes) / (1024.0 * 1024.0)) + << " repack_mb=" + << (static_cast(profile_ptr->repack_bytes) / (1024.0 * 1024.0)) + << " repack_calls=" << profile_ptr->repack_calls << '\n'; +} + +bool Smollm3TriAttentionKvCache::should_emit_keep_dump( + const Smollm3TriAttentionCompactionProfile* profile_ptr) const { + if (config_.dump_keep_path.empty()) + return false; + const int32_t dump_compaction_index = config_.dump_compaction_index; + if (dump_compaction_index <= 0) + return true; + return dump_compaction_index == (compaction_count_ + (profile_ptr == nullptr ? 1 : 0)); +} + +void Smollm3TriAttentionKvCache::emit_keep_dump_json( + const std::vector& keep_indices, int32_t keep_count, + const std::vector>& new_positions_by_head, + const Smollm3TriAttentionCompactionProfile* profile_ptr, + const std::vector& score_cache_files, + const std::vector& value_cache_files, + const std::vector& post_score_cache_files, + const std::vector& post_value_cache_files) { + if (!should_emit_keep_dump(profile_ptr)) + return; + json dump; + dump["compaction_index"] = profile_ptr != nullptr ? compaction_count_ : (compaction_count_ + 1); + dump["absolute_position"] = absolute_position_; + dump["cache_length_before"] = static_cast(cache_positions_.size()); + dump["keep_count"] = keep_count; + dump["prompt_end_position"] = prompt_end_position_; + dump["planned_prompt_length"] = planned_prompt_length_; + dump["protect_prefill"] = config_.protect_prefill; + dump["recent_window"] = config_.recent_window; + dump["kv_budget"] = config_.kv_budget; + dump["count_prompt_tokens"] = config_.count_prompt_tokens; + dump["cache_positions"] = cache_positions_; + dump["cache_positions_per_head"] = cache_positions_per_head_; + dump["new_positions_by_head"] = new_positions_by_head; + std::vector> keep_by_head(static_cast(cache_head_count_)); + for (int32_t cache_head = 0; cache_head < cache_head_count_; ++cache_head) { + auto begin = keep_indices.begin() + static_cast(cache_head * keep_count); + keep_by_head[static_cast(cache_head)] = + std::vector(begin, begin + keep_count); + } + dump["keep_indices_by_head"] = keep_by_head; + if (profile_ptr != nullptr) { + dump["profile"] = { + {"reserved_count", profile_ptr->reserved_count}, + {"candidate_count", profile_ptr->candidate_count}, + {"sampled_layers", profile_ptr->sampled_layers}, + {"sampled_heads", profile_ptr->sampled_heads}, + {"score_ms", profile_ptr->score_ms}, + {"combine_ms", profile_ptr->combine_ms}, + {"select_ms", profile_ptr->select_ms}, + {"repack_ms", profile_ptr->repack_ms}, + }; + } + if (config_.dump_score_cache) { + dump["score_cache_shape"] = {static_cast(cache_positions_.size()), + cache_head_count_, stats_.head_dim}; + dump["score_cache_dtype"] = "float32"; + dump["score_cache_files"] = score_cache_files; + dump["value_cache_files"] = value_cache_files; + dump["post_score_cache_shape"] = {keep_count, cache_head_count_, stats_.head_dim}; + dump["post_score_cache_files"] = post_score_cache_files; + dump["post_value_cache_files"] = post_value_cache_files; + } + std::ofstream out(config_.dump_keep_path); + out << dump.dump(2); + out << '\n'; + if (config_.abort_after_dump) + throw std::runtime_error("TriAttention aborted after keep dump"); +} + +void Smollm3TriAttentionKvCache::advance(int32_t n_tokens) { + assert(n_tokens == 1 && "Smollm3TriAttentionKvCache::advance only supports n_tokens==1"); + (void)n_tokens; + + if (cache_length_ >= max_length_) + compact_existing_cache(true); + + const auto row_bytes = static_cast(kv_dim_) * cache_element_size_; + const auto offset = static_cast(cache_length_) * row_bytes; + for (int32_t i = 0; i < num_layers_; ++i) { + const auto li = static_cast(i); + cudaMemcpyAsync(static_cast(cache_k_[li].data()) + offset, present_k_[li].data(), + row_bytes, cudaMemcpyDeviceToDevice, stream_); + cudaMemcpyAsync(static_cast(cache_v_[li].data()) + offset, present_v_[li].data(), + row_bytes, cudaMemcpyDeviceToDevice, stream_); + } + + if (cache_length_ < max_length_) + ++cache_length_; + for (auto& head_positions : cache_positions_per_head_) + head_positions.push_back(absolute_position_); + sync_shared_positions_from_head0(); + ++absolute_position_; + + if (cache_length_ >= compaction_trigger_length()) + compact_existing_cache(); +} + +void Smollm3TriAttentionKvCache::set_prompt_length(int32_t prompt_length) { + planned_prompt_length_ = std::max(prompt_length, 0); +} + +void Smollm3TriAttentionKvCache::mark_prefill_complete() { + prompt_end_position_ = std::max(absolute_position_, planned_prompt_length_); +} + +void Smollm3TriAttentionKvCache::reset() { + cache_length_ = 0; + absolute_position_ = 0; + planned_prompt_length_ = 0; + prompt_end_position_ = 0; + compaction_count_ = 0; + cache_positions_.clear(); + for (auto& head_positions : cache_positions_per_head_) + head_positions.clear(); + // Reset only logical sequence metadata. Cache length and position maps ensure + // stale device rows are neither scored nor exposed to the attention engine. +} + +std::size_t Smollm3TriAttentionKvCache::device_memory_bytes() const { + std::size_t total = 0; + for (const auto& t : cache_k_) + total += t.nbytes(); + for (const auto& t : cache_v_) + total += t.nbytes(); + for (const auto& t : present_k_) + total += t.nbytes(); + for (const auto& t : present_v_) + total += t.nbytes(); +#ifdef TRTMC_HAS_CUDA_KERNELS + total += candidate_indices_device_.nbytes(); + total += keep_indices_device_.nbytes(); + total += positions_device_.nbytes(); + total += inv_freq_device_.nbytes(); + total += cos_phase_device_.nbytes(); + total += sin_phase_device_.nbytes(); + total += scratch_k_device_.nbytes(); + total += scratch_v_device_.nbytes(); + for (const auto& layer : layer_gpu_stats_) { + total += layer.head_offsets.nbytes(); + total += layer.head_cache_indices.nbytes(); + total += layer.q_mean_real.nbytes(); + total += layer.q_mean_imag.nbytes(); + total += layer.q_abs_mean.nbytes(); + total += layer.freq_scale_sq.nbytes(); + total += layer.scores.nbytes(); + } +#endif + return total; +} + +bool Smollm3TriAttentionKvCache::ok() const { + if (cache_k_.size() != static_cast(num_layers_)) + return false; + for (const auto& t : cache_k_) { + if (!t.ok()) + return false; + } + for (const auto& t : cache_v_) { + if (!t.ok()) + return false; + } + for (const auto& t : present_k_) { + if (!t.ok()) + return false; + } + for (const auto& t : present_v_) { + if (!t.ok()) + return false; + } + return true; +} + +} // namespace trtmc diff --git a/src/runtime/models/smollm3/triattention_kv_cache.h b/src/runtime/models/smollm3/triattention_kv_cache.h new file mode 100644 index 0000000000..0fc998d76a --- /dev/null +++ b/src/runtime/models/smollm3/triattention_kv_cache.h @@ -0,0 +1,380 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "runtime/models/smollm3/inference_state.h" +#include "runtime/models/smollm3/kv_cache.h" +#include "trtmc/runtime/device_tensor.h" + +#include +#include +#include +#include + +namespace trtmc { + +enum class Smollm3TriAttentionScoreAggregation { + kMean, + kMax, +}; + +enum class Smollm3TriAttentionRopeStyle { + kHalf, + kInterleaved, +}; + +struct Smollm3TriAttentionConfig { + bool enabled{false}; + int32_t kv_budget{0}; + int32_t divide_length{128}; + int32_t recent_window{128}; + Smollm3TriAttentionScoreAggregation score_aggregation{ + Smollm3TriAttentionScoreAggregation::kMean}; + Smollm3TriAttentionScoreAggregation per_layer_aggregation{ + Smollm3TriAttentionScoreAggregation::kMean}; + bool count_prompt_tokens{true}; + bool protect_prefill{true}; + bool disable_mlr{false}; + bool disable_trig{false}; + std::string stats_section{"triattention_stats.json"}; + int32_t offset_max_length{65536}; + + // Debug / profile knobs — previously read ad-hoc via std::getenv + // (TRTMC_TRIATTN_DEBUG, TRTMC_TRIATTN_PROFILE, TRTMC_TRIATTN_DISABLE_GPU_*, + // TRTMC_TRIATTN_DUMP_*, TRTMC_TRIATTN_ZERO_TAIL, + // TRTMC_TRIATTN_RUNTIME_BUCKET_ROWS). Now populated once at + // construction from the registry-supplied ConfigBundle; the runtime + // reads these struct fields instead of hitting the environment. + bool debug{false}; + bool profile{false}; + int32_t runtime_bucket_rows{32}; + bool disable_gpu_selection{false}; + bool disable_gpu_compaction{false}; + bool disable_gpu_state{false}; + bool zero_tail{false}; + std::string dump_keep_path{}; + int32_t dump_compaction_index{0}; + bool abort_after_dump{false}; + bool dump_score_cache{false}; + bool dump_score_values{false}; +}; + +struct Smollm3TriAttentionHeadStats { + std::vector q_mean_real; + std::vector q_mean_imag; + std::vector q_abs_mean; + std::vector freq_scale_sq; +}; + +struct Smollm3TriAttentionStats { + int32_t head_dim{0}; + Smollm3TriAttentionRopeStyle rope_style{Smollm3TriAttentionRopeStyle::kHalf}; + float rope_theta{10000.0F}; + int32_t num_attention_heads{0}; + int32_t num_key_value_heads{0}; + int32_t stats_head_count{0}; + int32_t num_layers{0}; + std::vector inv_freq; + std::vector> sampled_score_heads_by_layer; + std::vector layer_stats; +}; + +struct Smollm3TriAttentionCompactionProfile { + double host_copy_ms{0.0}; + double host_convert_ms{0.0}; + double trig_prep_ms{0.0}; + double score_ms{0.0}; + double combine_ms{0.0}; + double select_ms{0.0}; + double repack_ms{0.0}; + std::size_t host_copy_bytes{0}; + std::size_t repack_bytes{0}; + int64_t repack_calls{0}; + int32_t candidate_count{0}; + int32_t reserved_count{0}; + int32_t sampled_layers{0}; + int32_t sampled_heads{0}; +}; + +// Forward declaration — full type in trtmc/config/config_bundle.h. +namespace config { +class ConfigBundle; +} + +// Build a Smollm3TriAttentionConfig from the bundle's JSON plus (optionally) the +// session-resolved ConfigBundle. Legacy bundles without the generic +// `defaults:` block still parse through the JSON path; when ``runtime_config`` +// has a non-default layer value for a field, that value wins. Environment +// variables (TRTMC_TRIATTN_*) are no longer read — callers supply values via +// the registry (CLI --set / --config or bundle defaults:). +Smollm3TriAttentionConfig +smollm3_parse_triattention_bundle_config(const std::string& config_json, int32_t max_cache_length, + const config::ConfigBundle* runtime_config = nullptr); + +Smollm3TriAttentionStats smollm3_parse_triattention_stats_json(const std::string& stats_json, + int32_t num_attention_heads, + int32_t num_key_value_heads, + int32_t num_layers); + +class Smollm3TriAttentionKvCache : public Smollm3InferenceState { + public: + Smollm3TriAttentionKvCache(int32_t num_layers, int32_t num_kv_heads, int32_t max_length, + int32_t kv_dim, cudaStream_t stream, + Smollm3TriAttentionConfig config, Smollm3TriAttentionStats stats, + DType cache_dtype = DType::kFloat32, Smollm3KvCacheNames names = {}); + + void reset() override; + void bind_to(TrtModule& module) override; + void prepare_step(TensorMap& inputs, int32_t seq_len = 1) override; + void advance(int32_t n_tokens = 1) override; + void set_prompt_length(int32_t prompt_length) override; + void mark_prefill_complete() override; + int32_t position() const override { return absolute_position_; } + int32_t max_length() const override { return max_length_; } + int32_t preferred_cache_rows() const override; + int32_t num_layers() const override { return num_layers_; } + bool needs_attention_mask() const override { return true; } + std::size_t device_memory_bytes() const override; + const char* state_type() const override { return "triattention_kv_cache"; } + bool ok() const override; + + void build_attention_mask(std::vector& mask) const; + + DeviceTensor& cache_k(int32_t layer) { return cache_k_[static_cast(layer)]; } + DeviceTensor& cache_v(int32_t layer) { return cache_v_[static_cast(layer)]; } + DeviceTensor& present_k(int32_t layer) { return present_k_[static_cast(layer)]; } + DeviceTensor& present_v(int32_t layer) { return present_v_[static_cast(layer)]; } + + int32_t active_length() const { return cache_length_; } + const std::vector& cache_positions() const { return cache_positions_; } + const std::vector>& cache_positions_per_head() const { + return cache_positions_per_head_; + } + int32_t prompt_end_position() const { return prompt_end_position_; } + + private: + void validate_shapes(); + void normalize_sampled_heads(); + void log_init_debug() const; + void allocate_layer_tensors(); + int32_t compaction_trigger_length() const; + int32_t compaction_keep_budget(int32_t total_tokens) const; + int32_t count_prefix_rows() const; + void compact_existing_cache(bool reserve_slot_for_append = false); + void dump_cache_rows_to_file(const DeviceTensor& tensor, int32_t rows, const std::string& path, + std::vector& out_files); + void maybe_dump_score_cache(int32_t rows, const char* k_pattern, const char* v_pattern, + std::vector& score_files, + std::vector& value_files); + bool compact_layer_on_gpu(int32_t layer, const std::vector& keep_indices, + int32_t keep_count, std::size_t row_bytes, int64_t& repack_calls, + std::size_t& repack_bytes); +#ifdef TRTMC_HAS_CUDA_KERNELS + bool gpu_compaction_upload_keep(int32_t keep_count, const std::vector& keep_indices); + bool gpu_compact_one_layer(int32_t layer, int32_t keep_count, std::size_t row_bytes); +#endif + void compact_layer_on_host(int32_t layer, const std::vector& keep_indices, + int32_t keep_count, std::size_t row_bytes, + std::size_t head_block_bytes, int32_t old_cache_length, + int64_t& repack_calls, std::size_t& repack_bytes); + void finalize_repack_profile(Smollm3TriAttentionCompactionProfile* profile_ptr, + cudaEvent_t repack_start, cudaEvent_t repack_stop, + int64_t repack_calls, std::size_t repack_bytes) const; + std::vector> + build_new_positions_by_head(const std::vector& keep_indices, int32_t keep_count) const; + void log_compact_debug(const std::vector& representative_positions, + const std::vector& keep_indices, int32_t keep_count) const; + std::vector collect_dropped_positions(const std::vector& keep_indices, + int32_t keep_count) const; + int32_t count_prefix_positions(const std::vector& representative_positions) const; + void log_compact_profile(const Smollm3TriAttentionCompactionProfile* profile_ptr, + int32_t keep_count) const; + bool should_emit_keep_dump(const Smollm3TriAttentionCompactionProfile* profile_ptr) const; + void emit_keep_dump_json(const std::vector& keep_indices, int32_t keep_count, + const std::vector>& new_positions_by_head, + const Smollm3TriAttentionCompactionProfile* profile_ptr, + const std::vector& score_cache_files, + const std::vector& value_cache_files, + const std::vector& post_score_cache_files, + const std::vector& post_value_cache_files); + std::vector + select_keep_indices(int32_t keep_budget, + Smollm3TriAttentionCompactionProfile* profile = nullptr); + std::vector build_reserve_mask(int32_t total_tokens, int32_t old_budget) const; + std::vector broadcast_indices_per_head(std::vector rows, + int32_t row_count) const; + std::vector + select_keep_indices_host(int32_t keep_budget, const std::vector& reserved, + const std::vector& candidates, + Smollm3TriAttentionCompactionProfile* profile = nullptr) const; + std::vector + broadcast_reserved_for_empty_budget(int32_t keep_budget, + const std::vector& reserved) const; + void precompute_trig_phases(std::vector>& cos_phase, + std::vector>& sin_phase, int32_t half_dim, + Smollm3TriAttentionCompactionProfile* profile) const; + bool layer_stats_shapes_valid(const Smollm3TriAttentionHeadStats& layer_stats, + int32_t half_dim) const; + void extract_k_rot(const float* row_ptr, int32_t head_offset, int32_t d, int32_t half_dim, + float& k_rot_real, float& k_rot_imag) const; + float reduce_trig_sums(const std::vector& trig_sums) const; + float score_one_row(const float* row_ptr, const Smollm3TriAttentionHeadStats& layer_stats, + std::size_t stats_base, int32_t head_offset, int32_t half_dim, + const std::vector>& cos_phase, + const std::vector>& sin_phase) const; + void score_rows_for_head(std::vector& scores, const std::vector& layer_cache, + const Smollm3TriAttentionHeadStats& layer_stats, int32_t score_head, + int32_t cache_head, int32_t half_dim, int32_t total_tokens, + const std::vector>& cos_phase, + const std::vector>& sin_phase) const; + void standardize_scores(std::vector& scores) const; + void accumulate_layer_fallback(const std::vector>& layer_scores, + std::vector& global_fallback_sum, + int32_t& global_fallback_count, int32_t total_tokens) const; + void reduce_group_into_aggregate(int32_t cache_head, int32_t total_tokens, + const std::vector>& layer_scores, + const std::vector& sampled_group, + std::vector& aggregate, bool first_layer_for_cache_head, + float* layer_dump) const; + void + accumulate_layer_to_aggregate(int32_t layer, int32_t total_tokens, + const std::vector>& layer_scores, + const std::vector>& sampled_by_cache_head, + std::vector>& aggregated_scores, + std::vector& contributing_layers_by_cache_head, + std::vector& layer_aggregate_dump) const; + std::vector compute_fallback_mean(const std::vector& global_fallback_sum, + int32_t global_fallback_count, + int32_t total_tokens) const; + void finalize_per_head_aggregate(std::vector>& aggregated_scores, + const std::vector& contributing_layers_by_cache_head, + const std::vector& global_fallback_mean) const; + void maybe_dump_score_values(const std::vector>& aggregated_scores, + const std::vector& layer_aggregate_dump, + int32_t total_tokens) const; + std::vector + build_keep_indices_per_head(const std::vector>& aggregated_scores, + const std::vector& reserved, + const std::vector& candidates, int32_t keep_budget, + int32_t need) const; + bool process_layer_for_host_selection(int32_t layer, int32_t half_dim, int32_t total_tokens, + const std::vector>& cos_phase, + const std::vector>& sin_phase, + std::vector>& aggregated_scores, + std::vector& global_fallback_sum, + int32_t& global_fallback_count, + std::vector& contributing_layers_by_cache_head, + std::vector& layer_aggregate_dump, + Smollm3TriAttentionCompactionProfile* profile) const; + std::vector + copy_cache_rows_to_host(const DeviceTensor& tensor, int32_t rows, + Smollm3TriAttentionCompactionProfile* profile = nullptr) const; + void sync_shared_positions_from_head0(); +#ifdef TRTMC_HAS_CUDA_KERNELS + struct LayerGpuStats { + DeviceTensor head_offsets; + DeviceTensor head_cache_indices; + DeviceTensor q_mean_real; + DeviceTensor q_mean_imag; + DeviceTensor q_abs_mean; + DeviceTensor freq_scale_sq; + DeviceTensor scores; + std::vector host_cache_head_indices; + int32_t score_head_count{0}; + }; + + void initialize_gpu_state(); + void allocate_core_selection_buffers(int32_t half_dim); + void build_layer_gpu_stats(int32_t layer, int32_t half_dim); + bool can_use_gpu_selection() const; + bool core_selection_buffers_ready() const; + static bool layer_gpu_stats_ready(const LayerGpuStats& layer); + enum class GpuLayerResult { kSkipped, kContributed, kFailed }; + GpuLayerResult process_layer_for_gpu_selection( + int32_t layer, int32_t total_tokens, int32_t num_offsets, + std::vector& aggregated_scores, std::vector& global_fallback_sum, + int32_t& global_fallback_count, std::vector& contributing_layers_by_cache_head, + Smollm3TriAttentionCompactionProfile* profile); + void standardize_score_rows(float* rows, int32_t num_rows, int32_t total_tokens) const; + void accumulate_flat_fallback(const float* rows, int32_t num_rows, int32_t total_tokens, + std::vector& fallback_sum, int32_t& fallback_count) const; + void aggregate_gpu_layer_into_cache_heads( + const std::vector& host_scores, const LayerGpuStats& gpu, int32_t total_tokens, + std::vector& aggregated_scores, + std::vector& contributing_layers_by_cache_head) const; + bool upload_candidate_indices_identity(int32_t total_tokens); + bool upload_gpu_trig_phases(int32_t num_offsets, int32_t half_dim, + Smollm3TriAttentionCompactionProfile* profile); + bool run_gpu_selection_over_layers(int32_t total_tokens, int32_t num_offsets, + std::vector& aggregated_scores, + std::vector& global_fallback_sum, + int32_t& global_fallback_count, + std::vector& contributing_layers_by_cache_head, + int32_t& contributing_layers, + Smollm3TriAttentionCompactionProfile* profile); + void + finalize_flat_per_head_aggregate(std::vector& aggregated_scores, + const std::vector& contributing_layers_by_cache_head, + const std::vector& global_fallback_mean, + int32_t total_tokens) const; + std::vector build_keep_from_flat_aggregate(const std::vector& aggregated_scores, + const std::vector& reserved, + const std::vector& candidates, + int32_t keep_budget, int32_t need, + int32_t total_tokens) const; + std::vector + select_keep_indices_gpu(int32_t keep_budget, const std::vector& reserved, + const std::vector& candidates, + Smollm3TriAttentionCompactionProfile* profile = nullptr); +#endif + + std::vector cache_k_; + std::vector cache_v_; + std::vector present_k_; + std::vector present_v_; + int32_t num_layers_{0}; + int32_t num_kv_heads_{0}; + int32_t query_head_count_{0}; + int32_t query_group_size_{0}; + int32_t cache_head_count_{0}; + int32_t score_group_size_{0}; + int32_t max_length_{0}; + int32_t kv_dim_{0}; + int32_t cache_length_{0}; + int32_t absolute_position_{0}; + int32_t planned_prompt_length_{0}; + int32_t prompt_end_position_{0}; + cudaStream_t stream_{nullptr}; + std::vector mask_buf_; + int32_t pos_buf_{0}; + bool has_position_input_{false}; + bool dynamic_binding_enabled_{false}; + int32_t bound_cache_rows_{0}; + DType cache_dtype_{DType::kFloat32}; + std::size_t cache_element_size_{sizeof(float)}; + Smollm3KvCacheNames names_; + Smollm3TriAttentionConfig config_; + Smollm3TriAttentionStats stats_; + std::vector cache_positions_; + std::vector> cache_positions_per_head_; + std::vector offsets_; + bool profile_enabled_{false}; + int64_t compaction_count_{0}; + TrtModule* bound_module_{nullptr}; +#ifdef TRTMC_HAS_CUDA_KERNELS + std::vector layer_gpu_stats_; + DeviceTensor candidate_indices_device_; + DeviceTensor keep_indices_device_; + DeviceTensor positions_device_; + DeviceTensor inv_freq_device_; + DeviceTensor cos_phase_device_; + DeviceTensor sin_phase_device_; + DeviceTensor scratch_k_device_; + DeviceTensor scratch_v_device_; +#endif +}; + +} // namespace trtmc diff --git a/tests/builder/test_native_kv_explicit_attention_contract.py b/tests/builder/test_native_kv_explicit_attention_contract.py index 1f8ba3f5b4..79d258ce69 100644 --- a/tests/builder/test_native_kv_explicit_attention_contract.py +++ b/tests/builder/test_native_kv_explicit_attention_contract.py @@ -14,6 +14,7 @@ _FIXED_KV_OWNERS = { _FAMILIES / "qwen" / "graph_ops.py", _FAMILIES / "llama" / "graph_ops.py", + _FAMILIES / "smollm3" / "graph_ops.py", _FAMILIES / "lfm2" / "model.py", _FAMILIES / "k2_horizon" / "model" / "model.py", } diff --git a/tests/cpp/models/smollm3/test_smollm3_chat_template.cpp b/tests/cpp/models/smollm3/test_smollm3_chat_template.cpp new file mode 100644 index 0000000000..9ed7ef281a --- /dev/null +++ b/tests/cpp/models/smollm3/test_smollm3_chat_template.cpp @@ -0,0 +1,205 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +// Model-owned chat-template coverage for smollm3 decoder formats. + +#include "runtime/models/smollm3/chat_templates.h" + +#include +#include + +static int failures = 0; + +static void check(bool condition, const char* test_name) { + if (!condition) { + std::cerr << "FAIL: " << test_name << '\n'; + ++failures; + } +} + +static void test_detect_nemotron_h() { + std::string tpl = "{% if add_generation_prompt %}System\n" + "User\n{{ message.content }}\n" + "Assistant\n{% endif %}"; + auto fmt = trtmc::smollm3_detect_chat_template_format(tpl); + check(fmt == "nemotron_h", "nemotron-h detection"); +} + +static void test_detect_chatml() { + std::string tpl = "{% for message in messages %}<|im_start|>{{ message.role }}\n{{ " + "message.content }}<|im_end|>\n{% endfor %}"; + auto fmt = trtmc::smollm3_detect_chat_template_format(tpl); + check(fmt == "chatml", "chatml detection"); +} + +static void test_detect_mistral() { + std::string tpl = "{{ bos_token }}{% for message in messages %}{% if message['role'] == 'user' " + "%}[INST] {{ message['content'] }} [/INST]{% endif %}{% endfor %}"; + auto fmt = trtmc::smollm3_detect_chat_template_format(tpl); + check(fmt == "mistral", "mistral detection"); +} + +static void test_detect_phi() { + std::string tpl = "{% for message in messages %}<|user|>\n{{ message.content " + "}}<|end|>\n<|assistant|>\n{% endfor %}"; + auto fmt = trtmc::smollm3_detect_chat_template_format(tpl); + check(fmt == "phi", "phi detection"); +} + +static void test_detect_gemma() { + std::string tpl = "{% for message in messages %}{{ message.role }}\n{{ " + "message.content }}\n{% endfor %}"; + auto fmt = trtmc::smollm3_detect_chat_template_format(tpl); + check(fmt == "gemma", "gemma detection"); +} + +static void test_detect_llama3() { + std::string tpl = "{% for message in messages %}<|start_header_id|>{{ message.role " + "}}<|end_header_id|>\n{{ message.content }}<|eot_id|>{% endfor %}"; + auto fmt = trtmc::smollm3_detect_chat_template_format(tpl); + check(fmt == "llama3", "llama3 detection"); +} + +static void test_apply_nemotron_h_no_thinking() { + auto result = trtmc::smollm3_apply_chat_template("nemotron_h", "hello", false); + check(result == "System\n\nUser\nhello\n" + "Assistant\n", + "nemotron-h no-thinking application"); +} + +static void test_apply_chatml_no_thinking() { + auto result = trtmc::smollm3_apply_chat_template("chatml", "What is 2+2?", false); + check(result == "<|im_start|>user\nWhat is " + "2+2?<|im_end|>\n<|im_start|>assistant\n\n\n\n\n", + "chatml no-thinking application"); +} + +static void test_apply_mistral_no_thinking_ignored() { + auto result = trtmc::smollm3_apply_chat_template("mistral", "hello", false); + check(result == "[INST] hello [/INST]", "mistral no-thinking ignored"); +} + +static void test_apply_phi() { + auto result = trtmc::smollm3_apply_chat_template("phi", "hello"); + check(result == "<|user|>\nhello<|end|>\n<|assistant|>\n", "phi application"); +} + +static void test_apply_gemma() { + auto result = trtmc::smollm3_apply_chat_template("gemma", "hello"); + check(result == "user\nhello\nmodel\n", + "gemma application"); +} + +static void test_apply_llama3() { + auto result = trtmc::smollm3_apply_chat_template("llama3", "hello"); + check(result == "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\nhello<|eot_id|><|" + "start_header_id|>assistant<|end_header_id|>\n\n", + "llama3 application"); +} + +// ── SmolLM3's own chat template ───────────────────────────────────────────── +// The expected strings below are the verbatim output of the upstream +// tokenizer's apply_chat_template() for HuggingFaceTB/SmolLM3-3B at revision +// a07cc9a0, captured with the date pinned. Asserting against the real +// rendering, rather than a reading of the Jinja source, is what makes this a +// contract: SmolLM3 always emits a system block, and it does *not* close that +// block with <|im_end|> before the user turn. + +static const char* kSmollm3Date = "04 September 2026"; + +static void test_detect_smollm3_over_chatml() { + // SmolLM3's template is ChatML-framed, so generic ChatML detection must not + // claim it -- the SmolLM3 branch carries the mandatory system block. + std::string tpl = "{%- if enable_thinking %}{%- set reasoning_mode = \"/think\" %}{%- endif %}" + "{{- \"<|im_start|>system\\n\" -}}" + "{{- \"Reasoning Mode: \" + reasoning_mode + \"\\n\\n\" -}}"; + check(trtmc::smollm3_detect_chat_template_format(tpl) == "smollm3", + "smollm3 detection wins over chatml"); +} + +static void test_apply_smollm3_no_thinking() { + const std::string expected = + "<|im_start|>system\n" + "## Metadata\n" + "\n" + "Knowledge Cutoff Date: June 2025\n" + "Today Date: 04 September 2026\n" + "Reasoning Mode: /no_think\n" + "\n" + "## Custom Instructions\n" + "\n" + "You are a helpful AI assistant named SmolLM, trained by Hugging Face.\n" + "\n" + "<|im_start|>user\n" + "What is 2+2?<|im_end|>\n" + "<|im_start|>assistant\n" + "\n" + "\n" + "\n"; + auto result = + trtmc::smollm3_apply_chat_template("smollm3", "What is 2+2?", false, kSmollm3Date); + check(result == expected, "smollm3 /no_think matches upstream byte for byte"); + check(result.size() == 297, "smollm3 /no_think length matches upstream (297)"); +} + +static void test_apply_smollm3_thinking() { + const std::string expected = + "<|im_start|>system\n" + "## Metadata\n" + "\n" + "Knowledge Cutoff Date: June 2025\n" + "Today Date: 04 September 2026\n" + "Reasoning Mode: /think\n" + "\n" + "## Custom Instructions\n" + "\n" + "You are a helpful AI assistant named SmolLM, trained by Hugging Face. Your role as an " + "assistant involves thoroughly exploring questions through a systematic thinking process " + "before providing the final precise and accurate solutions. This requires engaging in a " + "comprehensive cycle of analysis, summarizing, exploration, reassessment, reflection, " + "backtracking, and iteration to develop well-considered thinking process. Please structure " + "your response into two main sections: Thought and Solution using the specified format: " + " Thought section Solution section. In the Thought section, detail your " + "reasoning process in steps. Each step should include detailed considerations such as " + "analysing questions, summarizing relevant findings, brainstorming new ideas, verifying " + "the accuracy of the current steps, refining any errors, and revisiting previous steps. In " + "the Solution section, based on various attempts, explorations, and reflections from the " + "Thought section, systematically present the final solution that you deem correct. The " + "Solution section should be logical, accurate, and concise and detail necessary steps " + "needed to reach the conclusion.\n" + "\n" + "<|im_start|>user\n" + "What is 2+2?<|im_end|>\n" + "<|im_start|>assistant\n"; + auto result = trtmc::smollm3_apply_chat_template("smollm3", "What is 2+2?", true, kSmollm3Date); + check(result == expected, "smollm3 /think matches upstream byte for byte"); + check(result.size() == 1369, "smollm3 /think length matches upstream (1369)"); +} + +int main() { + + test_detect_smollm3_over_chatml(); + test_detect_chatml(); + test_detect_mistral(); + test_detect_phi(); + test_detect_gemma(); + test_detect_llama3(); + test_detect_nemotron_h(); + test_apply_smollm3_no_thinking(); + test_apply_smollm3_thinking(); + test_apply_chatml_no_thinking(); + test_apply_mistral_no_thinking_ignored(); + test_apply_phi(); + test_apply_gemma(); + test_apply_llama3(); + test_apply_nemotron_h_no_thinking(); + + if (failures > 0) { + std::cerr << failures << " test(s) FAILED\n"; + return 1; + } + std::cerr << "All smollm3 chat_template tests passed.\n"; + return 0; +} diff --git a/tests/cpp/models/smollm3/test_smollm3_native_kv_cache.cpp b/tests/cpp/models/smollm3/test_smollm3_native_kv_cache.cpp new file mode 100644 index 0000000000..6e294feb45 --- /dev/null +++ b/tests/cpp/models/smollm3/test_smollm3_native_kv_cache.cpp @@ -0,0 +1,54 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../../native_kv_cache_contract_test.h" +#include "runtime/models/smollm3/kv_cache.h" +#include "runtime/models/smollm3/pipeline.h" +#include "runtime/models/smollm3/plugin_helpers.h" + +namespace { + +int test_dynamic_prefill_uses_runtime_cache_rows() { + cudaStream_t stream = nullptr; + if (cudaStreamCreate(&stream) != cudaSuccess) + return 1; + + int failures = 0; + { + trtmc::Smollm3KvCache cache(1, 1000, 2, stream, trtmc::DType::kFloat16); + trtmc::test::NativeKvModuleStub prefill(stream, 1, 131072, 1, 2, trtmc::DType::kFloat16, + /*native=*/false, nullptr, 4, 16, + /*dynamic_legacy_cache=*/true); + + cache.bind_cache_inputs(prefill); + if (!trtmc::cache_input_supports_runtime_rows(prefill, "cache_k_0")) { + std::cerr << "FAIL [SmolLM3]: dynamic metadata survives positive active tensor shape\n"; + ++failures; + } + if (prefill.bound_shape("cache_k_0") != std::vector{1000, 2} || + prefill.bound_shape("cache_v_0") != std::vector{1000, 2}) { + std::cerr << "FAIL [SmolLM3]: dynamic prefill binds runtime-sized cache rows\n"; + ++failures; + } + + trtmc::TensorMap inputs; + cache.prepare_step(inputs, 4); + if (inputs.at("attention_mask").shape != std::vector{4, 1004}) { + std::cerr << "FAIL [SmolLM3]: dynamic prefill mask uses runtime cache rows\n"; + ++failures; + } + } + cudaStreamDestroy(stream); + return failures; +} + +} // namespace + +int main() { + return trtmc::test::run_native_kv_contract_tests("SmolLM3") + + test_dynamic_prefill_uses_runtime_cache_rows(); +} diff --git a/tests/cpp/models/smollm3/test_smollm3_pipeline.cpp b/tests/cpp/models/smollm3/test_smollm3_pipeline.cpp new file mode 100644 index 0000000000..bc5907a0bf --- /dev/null +++ b/tests/cpp/models/smollm3/test_smollm3_pipeline.cpp @@ -0,0 +1,512 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +// ============================================================================= +// ISO 26262 Traceability +// ============================================================================= +// Trace ID: UT-DEC-CPP-02 +// Architecture: ARCH-FAC-001 +// Unit Design: UD-TRT-DEC-01 +// Intent: Smollm3TextGenerationPipeline prefill/decode loop, argmax selection, EOS stopping +// Preconditions: TRT + CUDA GPU available, identity engine built in-process +// Postconditions: Pipeline generates correct tokens, stops at EOS, respects max_new_tokens +// ============================================================================= + +// ============================================================================= +// Test suite: SmolLM3-owned Smollm3TextGenerationPipeline copy +// ============================================================================= +// +// Tests the Smollm3TextGenerationPipeline using a tiny TRT identity engine. +// The identity engine maps token_id[1] → logits[4] (just copies input to output). +// This validates the prefill→decode loop, argmax, and EOS stopping. +// +// For full E2E validation with real models, see tests/test_e2e.py. +// ============================================================================= + +#include "runtime/models/smollm3/kv_cache.h" +#include "runtime/models/smollm3/pipeline.h" +#include "trtmc/runtime/trt_module.h" +#include "trtmc/tokenizer.h" +// pipeline_interface.h was removed; GenerateConfig is in trtmc/pipeline.h +// (already included transitively via runtime/models/smollm3/pipeline.h) + +#include "runtime/backend/trt_module_impl.h" +#include "runtime/core/trt_common.h" + +#include +#include +#include +#include +#include +#include + +static int failures = 0; + +static void check(bool condition, const char* test_name) { + if (!condition) { + std::cerr << "FAIL: " << test_name << '\n'; + ++failures; + } +} + +static trtmc::TrtLogger g_logger; + +namespace trtmc { + +class TrtModuleImplTestPeer { + public: + static void set_execution_context_name(TrtModuleImpl& module, const char* name) { + module.ctx_->setName(name); + } + + static std::string execution_context_name(const TrtModuleImpl& module) { + return module.ctx_->getName(); + } +}; + +} // namespace trtmc + +class MockTokenizer final : public trtmc::ITokenizer { + public: + std::vector encode(const std::string& text) const override { + (void)text; + return {9}; + } + + std::string decode(const std::vector& ids) const override { + std::string out; + for (int32_t id : ids) { + out += token_for_id(id); + } + return out; + } + + int32_t id_for_token(std::string_view token) const override { + if (token == "\\boxed{") + return 1; + if (token == "70") + return 2; + if (token == "}") + return 3; + if (token == " extra") + return 4; + return 0; + } + + std::string token_for_id(int32_t id) const override { + switch (id) { + case 1: + return "\\boxed{"; + case 2: + return "70"; + case 3: + return "}"; + case 4: + return " extra"; + default: + return ""; + } + } +}; + +class SequenceSampler final : public trtmc::Smollm3ISampler { + public: + explicit SequenceSampler(std::vector tokens) : tokens_(std::move(tokens)) {} + + trtmc::Smollm3SampleResult sample(const float* logits, int32_t vocab_size, + const trtmc::Smollm3SamplingParams& params) override { + (void)logits; + (void)vocab_size; + trtmc::Smollm3SampleResult result; + const std::size_t idx = cursor_ < tokens_.size() ? cursor_ : (tokens_.size() - 1); + result.token_id = tokens_[idx]; + result.is_eos = (result.token_id == params.eos_token_id); + if (cursor_ < tokens_.size()) + ++cursor_; + return result; + } + + trtmc::Smollm3LogitsLocation logits_location() const override { + return trtmc::Smollm3LogitsLocation::HOST; + } + const char* sampler_type() const override { return "sequence"; } + void reset() override { cursor_ = 0; } + + private: + std::vector tokens_; + std::size_t cursor_{0}; +}; + +// Build a tiny decoder-like engine: +// Inputs: token_id [1] int32, attention_mask [8] float32 +// Outputs: logits [4] float32 +// The engine produces fixed logits [0.1, 0.2, 0.9, 0.3] regardless of input +// (identity on a constant), so argmax always returns 2. +static trtmc::TrtUniquePtr build_mock_decoder() { + auto builder = trtmc::TrtUniquePtr(nvinfer1::createInferBuilder(g_logger)); + if (!builder) + return nullptr; + + auto network = trtmc::TrtUniquePtr(builder->createNetworkV2(0)); + auto config = trtmc::TrtUniquePtr(builder->createBuilderConfig()); + config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, 1 << 20); + + // Inputs + auto* token_inp = + network->addInput("token_id", nvinfer1::DataType::kINT32, nvinfer1::Dims{1, {1}}); + auto* mask_inp = + network->addInput("attention_mask", nvinfer1::DataType::kFLOAT, nvinfer1::Dims{1, {8}}); + + // Constant logits: [0.1, 0.2, 0.9, 0.3] — argmax = index 2 + float const_logits[4] = {0.1f, 0.2f, 0.9f, 0.3f}; + auto* const_w = network->addConstant( + nvinfer1::Dims{1, {4}}, nvinfer1::Weights{nvinfer1::DataType::kFLOAT, const_logits, 4}); + if (!const_w) + return nullptr; + + auto* out = const_w->getOutput(0); + out->setName("logits"); + network->markOutput(*out); + + // Need to "use" the inputs so TRT doesn't optimize them away + // Add identity on token_id and mask (mark as outputs too, then unmark) + // Actually, for a proper test engine, just mark them as used via identity + auto* id_token = network->addIdentity(*token_inp); + id_token->getOutput(0)->setName("_unused_token"); + + auto* id_mask = network->addIdentity(*mask_inp); + id_mask->getOutput(0)->setName("_unused_mask"); + + auto plan = trtmc::TrtUniquePtr( + builder->buildSerializedNetwork(*network, *config)); + if (!plan) + return nullptr; + + auto runtime = trtmc::TrtUniquePtr(nvinfer1::createInferRuntime(g_logger)); + return trtmc::TrtUniquePtr( + runtime->deserializeCudaEngine(plan->data(), plan->size())); +} + +static void test_pipeline_construction() { + auto engine = build_mock_decoder(); + if (!engine) { + std::cerr << "WARNING: Could not build mock decoder engine, skipping test\n"; + return; + } + + cudaStream_t stream; + cudaStreamCreate(&stream); + + auto module = std::make_unique(engine.get(), + engine->createExecutionContext(), stream); + auto cache = std::make_unique(1, 8, 4, stream); + + trtmc::Smollm3TextGenConfig cfg; + cfg.vocab_size = 4; + cfg.id_bos = 0; + cfg.id_eos = 2; // argmax will always hit this! + cfg.has_position_input = false; + + trtmc::Smollm3TextGenerationPipeline pipeline(std::move(module), std::move(cache), cfg, stream); + + check(std::string(pipeline.pipeline_type()) == "Smollm3TextGenerationPipeline", + "pipeline name"); + + cudaStreamDestroy(stream); +} + +static void test_generate_stops_at_eos() { + auto engine = build_mock_decoder(); + if (!engine) { + std::cerr << "WARNING: Could not build mock decoder engine, skipping test\n"; + return; + } + + cudaStream_t stream; + cudaStreamCreate(&stream); + + auto module = std::make_unique(engine.get(), + engine->createExecutionContext(), stream); + auto cache = std::make_unique(1, 8, 4, stream); + + trtmc::Smollm3TextGenConfig cfg; + cfg.vocab_size = 4; + cfg.id_bos = 0; + cfg.id_eos = 2; // argmax of [0.1, 0.2, 0.9, 0.3] = 2 = eos + cfg.has_position_input = false; + + trtmc::Smollm3TextGenerationPipeline pipeline(std::move(module), std::move(cache), cfg, stream); + + trtmc::GenerateConfig gen_cfg; + gen_cfg.max_new_tokens = 10; + + auto result = pipeline.generate_ids({1}, gen_cfg); + + // Input [1] + one generated token (eos=2) → should stop immediately + check(result.token_ids.size() == 2, "output has 2 tokens (input + eos)"); + check(result.token_ids[0] == 1, "first token is input"); + check(result.token_ids[1] == 2, "second token is eos (argmax=2)"); + + cudaStreamDestroy(stream); +} + +static void test_generate_stops_at_any_default_eos() { + auto engine = build_mock_decoder(); + if (!engine) + return; + + cudaStream_t stream; + cudaStreamCreate(&stream); + + auto module = std::make_unique(engine.get(), + engine->createExecutionContext(), stream); + auto cache = std::make_unique(1, 8, 4, stream); + + trtmc::Smollm3TextGenConfig cfg; + cfg.vocab_size = 4; + cfg.id_eos = 1; + cfg.id_eos_ids = {1, 2}; + cfg.has_position_input = false; + + trtmc::Smollm3TextGenerationPipeline pipeline(std::move(module), std::move(cache), cfg, stream); + + trtmc::GenerateConfig gen_cfg; + gen_cfg.max_new_tokens = 10; + + auto result = pipeline.generate_ids({1}, gen_cfg); + check(result.token_ids == std::vector({1, 2}), + "second configured EOS stops generation"); + + cudaStreamDestroy(stream); +} + +static void test_explicit_eos_override_replaces_default_set() { + auto engine = build_mock_decoder(); + if (!engine) + return; + + cudaStream_t stream; + cudaStreamCreate(&stream); + + auto module = std::make_unique(engine.get(), + engine->createExecutionContext(), stream); + auto cache = std::make_unique(1, 8, 4, stream); + + trtmc::Smollm3TextGenConfig cfg; + cfg.vocab_size = 4; + cfg.id_eos = 1; + cfg.id_eos_ids = {1, 2}; + cfg.has_position_input = false; + + trtmc::Smollm3TextGenerationPipeline pipeline(std::move(module), std::move(cache), cfg, stream); + + trtmc::GenerateConfig gen_cfg; + gen_cfg.max_new_tokens = 3; + gen_cfg.eos_token_id = 3; + + auto result = pipeline.generate_ids({1}, gen_cfg); + check(result.token_ids == std::vector({1, 2, 2, 2}), + "explicit scalar EOS override replaces the default EOS set"); + + cudaStreamDestroy(stream); +} + +static void test_generate_max_tokens() { + auto engine = build_mock_decoder(); + if (!engine) { + std::cerr << "WARNING: Could not build mock decoder engine, skipping test\n"; + return; + } + + cudaStream_t stream; + cudaStreamCreate(&stream); + + auto module = std::make_unique(engine.get(), + engine->createExecutionContext(), stream); + auto cache = std::make_unique(1, 8, 4, stream); + + trtmc::Smollm3TextGenConfig cfg; + cfg.vocab_size = 4; + cfg.id_bos = 0; + cfg.id_eos = 99; // EOS token that argmax will never produce + cfg.has_position_input = false; + + trtmc::Smollm3TextGenerationPipeline pipeline(std::move(module), std::move(cache), cfg, stream); + + trtmc::GenerateConfig gen_cfg; + gen_cfg.max_new_tokens = 3; + + auto result = pipeline.generate_ids({1}, gen_cfg); + + // Input [1] + 3 generated tokens (all argmax=2, never hits eos=99) + check(result.token_ids.size() == 4, "output has 4 tokens (input + 3 generated)"); + check(result.token_ids[0] == 1, "first = input"); + check(result.token_ids[1] == 2, "gen 1 = argmax(2)"); + check(result.token_ids[2] == 2, "gen 2 = argmax(2)"); + check(result.token_ids[3] == 2, "gen 3 = argmax(2)"); + + cudaStreamDestroy(stream); +} + +static void test_argmax() { + std::vector logits = {0.1f, 0.5f, 0.3f, 0.8f, 0.2f}; + int32_t result = trtmc::Smollm3TextGenerationPipeline::argmax(logits); + check(result == 3, "argmax of [0.1, 0.5, 0.3, 0.8, 0.2] = 3"); + + std::vector single = {42.0f}; + check(trtmc::Smollm3TextGenerationPipeline::argmax(single) == 0, "argmax of single = 0"); + + std::vector empty; + check(trtmc::Smollm3TextGenerationPipeline::argmax(empty) == 0, "argmax of empty = 0"); +} + +static void test_zero_max_tokens() { + auto engine = build_mock_decoder(); + if (!engine) + return; + + cudaStream_t stream; + cudaStreamCreate(&stream); + + auto module = std::make_unique(engine.get(), + engine->createExecutionContext(), stream); + auto cache = std::make_unique(1, 8, 4, stream); + + trtmc::Smollm3TextGenConfig cfg; + cfg.vocab_size = 4; + cfg.id_eos = 2; + cfg.has_position_input = false; + + trtmc::Smollm3TextGenerationPipeline pipeline(std::move(module), std::move(cache), cfg, stream); + + trtmc::GenerateConfig gen_cfg; + gen_cfg.max_new_tokens = 0; + + auto result = pipeline.generate_ids({1, 2, 3}, gen_cfg); + check(result.token_ids.size() == 3, "zero max_new_tokens returns input unchanged"); + + cudaStreamDestroy(stream); +} + +static void test_kv_reset_is_logical_and_masks_stale_rows() { + cudaStream_t stream; + cudaStreamCreate(&stream); + + trtmc::Smollm3KvCache cache(1, 8, 4, stream); + std::vector stale_k(32, 3.25F); + std::vector stale_v(32, -7.5F); + check(cache.cache_k(0).copy_from_host(stale_k.data()), "upload stale K cache rows"); + check(cache.cache_v(0).copy_from_host(stale_v.data()), "upload stale V cache rows"); + cache.set_position(5); + + cache.reset(); + + std::vector actual_k(stale_k.size()); + std::vector actual_v(stale_v.size()); + check(cache.cache_k(0).copy_to_host(actual_k.data()), "download stale K cache rows"); + check(cache.cache_v(0).copy_to_host(actual_v.data()), "download stale V cache rows"); + check(actual_k == stale_k, "logical reset preserves allocated K cache storage"); + check(actual_v == stale_v, "logical reset preserves allocated V cache storage"); + check(cache.position() == 0, "logical reset clears the visible cache length"); + + trtmc::TensorMap inputs; + cache.prepare_step(inputs); + const auto mask_it = inputs.find("attention_mask"); + check(mask_it != inputs.end(), "logical reset creates an attention mask"); + if (mask_it != inputs.end()) { + const auto& mask = mask_it->second; + check(mask.shape == std::vector{9}, "logical reset mask covers cache and token"); + const auto* values = static_cast(mask.data); + bool stale_rows_hidden = values != nullptr; + for (int32_t i = 0; stale_rows_hidden && i < 8; ++i) + stale_rows_hidden = values[i] < -1000.0F; + check(stale_rows_hidden, "logical reset masks every stale cache row"); + check(values != nullptr && values[8] == 0.0F, + "logical reset keeps the current token visible"); + } + + cudaStreamDestroy(stream); +} + +static void test_generation_reset_reuses_execution_context() { + auto engine = build_mock_decoder(); + if (!engine) + return; + + cudaStream_t stream; + cudaStreamCreate(&stream); + + auto* ctx = engine->createExecutionContext(); + trtmc::TrtModuleImpl module(engine.get(), ctx, stream); + trtmc::TrtModuleImplTestPeer::set_execution_context_name(module, "generation-context"); + + module.reset_execution_context(); + + check(trtmc::TrtModuleImplTestPeer::execution_context_name(module) == "generation-context", + "generation reset reuses the loaded execution context"); + + int32_t token_id = 7; + std::vector attention_mask(8, 0.0F); + trtmc::TensorMap inputs; + inputs["token_id"] = trtmc::Tensor{&token_id, {1}, trtmc::DType::kInt32}; + inputs["attention_mask"] = trtmc::Tensor{attention_mask.data(), {8}, trtmc::DType::kFloat32}; + auto outputs = module.forward(inputs); + check(outputs.count("logits") == 1, "reused execution context remains executable"); + + cudaStreamDestroy(stream); +} + +static void test_stop_on_boxed_answer() { + auto engine = build_mock_decoder(); + if (!engine) + return; + + cudaStream_t stream; + cudaStreamCreate(&stream); + + auto module = std::make_unique(engine.get(), + engine->createExecutionContext(), stream); + auto cache = std::make_unique(1, 8, 4, stream); + auto tokenizer = std::make_shared(); + auto sampler = std::make_unique(std::vector{1, 2, 3, 4}); + + trtmc::Smollm3TextGenConfig cfg; + cfg.vocab_size = 4; + cfg.id_eos = 99; + cfg.has_position_input = false; + + trtmc::Smollm3TextGenerationPipeline pipeline(std::move(module), std::move(cache), cfg, stream, + tokenizer, "mock", std::move(sampler)); + + trtmc::GenerateConfig gen_cfg; + gen_cfg.max_new_tokens = 10; + gen_cfg.stop_on_boxed_answer = true; + gen_cfg.stop_check_interval = 1; + + auto result = pipeline.generate_ids({9}, gen_cfg); + check(result.token_ids.size() == 4, "boxed-answer stop truncates generation"); + check(result.token_ids[1] == 1, "boxed stop token 1"); + check(result.token_ids[2] == 2, "boxed stop token 2"); + check(result.token_ids[3] == 3, "boxed stop token 3"); + + cudaStreamDestroy(stream); +} + +int main() { + test_argmax(); + test_pipeline_construction(); + test_generate_stops_at_eos(); + test_generate_stops_at_any_default_eos(); + test_explicit_eos_override_replaces_default_set(); + test_generate_max_tokens(); + test_zero_max_tokens(); + test_kv_reset_is_logical_and_masks_stale_rows(); + test_generation_reset_reuses_execution_context(); + test_stop_on_boxed_answer(); + + if (failures > 0) + std::cerr << failures << " test(s) FAILED\n"; + return failures; +} diff --git a/tests/cpp/models/smollm3/test_smollm3_plugin_helpers.cpp b/tests/cpp/models/smollm3/test_smollm3_plugin_helpers.cpp new file mode 100644 index 0000000000..a496566171 --- /dev/null +++ b/tests/cpp/models/smollm3/test_smollm3_plugin_helpers.cpp @@ -0,0 +1,206 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +// Unit tests for runtime plugin helper parsing. +// Focus: tokenizer_add_special_tokens detection from bundle config. + +#include "../../native_kv_cache_contract_test.h" +#include "runtime/models/smollm3/plugin_helpers.h" + +#include +#include +#include + +static int failures = 0; + +static void check(bool condition, const char* name) { + if (!condition) { + std::cerr << "FAIL: " << name << '\n'; + ++failures; + } +} + +static trtmc::BundleFile make_bundle_with_config(const std::string& config) { + trtmc::BundleFile bundle; + trtmc::BundleSection sec; + sec.name = "config.json"; + sec.data.assign(config.begin(), config.end()); + bundle.sections.push_back(std::move(sec)); + return bundle; +} + +static trtmc::BundleFile make_bundle_with_config_and_tokenizer(const std::string& config, + const std::string& tokenizer_json) { + auto bundle = make_bundle_with_config(config); + trtmc::BundleSection tok; + tok.name = "tokenizer.json"; + tok.data.assign(tokenizer_json.begin(), tokenizer_json.end()); + bundle.sections.push_back(std::move(tok)); + return bundle; +} + +static void check_ids(const std::vector& actual, const std::vector& expected, + const char* name) { + if (actual != expected) { + std::cerr << "FAIL: " << name << '\n'; + ++failures; + } +} + +static void test_missing_field_defaults_true() { + auto bundle = make_bundle_with_config(R"({"runtime_strategy":"smollm3_decoder_kv_cache"})"); + check(trtmc::detect_add_special_tokens(bundle) == true, + "detect_add_special_tokens: missing field defaults true"); +} + +static void test_integer_false_parsed() { + auto bundle = make_bundle_with_config(R"({"tokenizer_add_special_tokens":0})"); + check(trtmc::detect_add_special_tokens(bundle) == false, + "detect_add_special_tokens: integer 0 parsed as false"); +} + +static void test_integer_true_parsed() { + auto bundle = make_bundle_with_config(R"({"tokenizer_add_special_tokens":1})"); + check(trtmc::detect_add_special_tokens(bundle) == true, + "detect_add_special_tokens: integer 1 parsed as true"); +} + +static void test_bool_false_parsed() { + auto bundle = make_bundle_with_config(R"({"tokenizer_add_special_tokens":false})"); + check(trtmc::detect_add_special_tokens(bundle) == false, + "detect_add_special_tokens: bool false parsed as false"); +} + +static void test_bool_true_parsed() { + auto bundle = make_bundle_with_config(R"({"tokenizer_add_special_tokens":true})"); + check(trtmc::detect_add_special_tokens(bundle) == true, + "detect_add_special_tokens: bool true parsed as true"); +} + +static void test_decoder_profile_selection_keeps_runtime_ceiling() { + check_ids(trtmc::select_decoder_profile_rows({256, 131072}, 1000), {256, 131072}, + "decoder profile selection keeps first runtime ceiling"); + check_ids(trtmc::select_decoder_profile_rows({256, 131072}, 256), {256}, + "decoder profile selection stops at exact runtime capacity"); + bool rejected = false; + try { + (void)trtmc::select_decoder_profile_rows({256}, 131072); + } catch (const std::runtime_error&) { + rejected = true; + } + check(rejected, "decoder profile selection rejects an undersized largest bucket"); +} + +static void test_dynamic_profile_rows_use_profile_metadata() { + trtmc::test::NativeKvModuleStub module(nullptr, 1, 131072, 1, 2, trtmc::DType::kFloat16, + /*native=*/false, nullptr, 4, 16, + /*dynamic_legacy_cache=*/true, {131072, 256, 131072}, + {131072, 1, 1}); + + check(module.tensor_shape("cache_k_0") == std::vector{131072, 2}, + "dynamic module reports a positive active cache shape"); + check(trtmc::cache_input_supports_runtime_rows(module, "cache_k_0"), + "dynamic input metadata enables runtime rows despite positive active shape"); + check(trtmc::decoder_profile_cache_rows(module, "cache_k_0", 1, 131072) == 256, + "first dynamic decode profile uses its own row ceiling"); + check(trtmc::decoder_profile_cache_rows(module, "cache_k_0", 2, 131072) == 131072, + "second dynamic decode profile uses its own row ceiling"); + + const auto roles = trtmc::detect_decoder_profile_roles(module, "token_id", "cache_k_0", 131072); + check(roles.prefill_profile_idx == 0 && roles.prefill_max_length == 131072, + "role detection keeps the dynamic prefill profile"); + check(roles.decode_profiles.size() == 2 && roles.decode_profiles[0].profile_idx == 1 && + roles.decode_profiles[0].kv_rows == 256 && + roles.decode_profiles[1].profile_idx == 2 && + roles.decode_profiles[1].kv_rows == 131072, + "role detection preserves per-profile dynamic KV ceilings"); +} + +static void test_exact_special_frame_overrides_native_unigram_fallback() { + const std::string tokenizer_json = R"({ + "model": { + "type": "Unigram", + "unk_id": 0, + "vocab": [ + ["", 0.0], + ["", 0.0], + ["", 0.0], + ["\u2581", -1.0], + ["h", -1.0], + ["e", -1.0] + ] + }, + "pre_tokenizer": { + "type": "Metaspace", + "replacement": "\u2581", + "add_prefix_space": true + } + })"; + auto bundle = make_bundle_with_config_and_tokenizer( + R"({ + "tokenizer_add_special_tokens": 1, + "tokenizer_special_prefix_ids": [1], + "tokenizer_special_suffix_ids": [] + })", + tokenizer_json); + + auto tokenizer = trtmc::create_tokenizer_from_bundle(bundle); + check(tokenizer != nullptr, "create native tokenizer with exact special frame"); + check_ids(tokenizer->encode("he"), {1, 3, 4, 5}, + "exact special frame adds BOS without fallback EOS"); +} + +static void test_exact_special_frame_respects_add_special_false() { + const std::string tokenizer_json = R"({ + "model": { + "type": "Unigram", + "unk_id": 0, + "vocab": [ + ["", 0.0], + ["", 0.0], + ["", 0.0], + ["\u2581", -1.0], + ["h", -1.0], + ["e", -1.0] + ] + }, + "pre_tokenizer": { + "type": "Metaspace", + "replacement": "\u2581", + "add_prefix_space": true + } + })"; + auto bundle = make_bundle_with_config_and_tokenizer( + R"({ + "tokenizer_add_special_tokens": 1, + "tokenizer_special_prefix_ids": [1], + "tokenizer_special_suffix_ids": [2] + })", + tokenizer_json); + + auto tokenizer = trtmc::try_create_native_tokenizer(bundle, /*add_special_tokens=*/false); + check(tokenizer != nullptr, "create native tokenizer with special frame disabled"); + check_ids(tokenizer->encode("he"), {3, 4, 5}, + "exact special frame does not override add_special_tokens=false"); +} + +int main() { + test_missing_field_defaults_true(); + test_integer_false_parsed(); + test_integer_true_parsed(); + test_bool_false_parsed(); + test_bool_true_parsed(); + test_decoder_profile_selection_keeps_runtime_ceiling(); + test_dynamic_profile_rows_use_profile_metadata(); + test_exact_special_frame_overrides_native_unigram_fallback(); + test_exact_special_frame_respects_add_special_false(); + + if (failures > 0) { + std::cerr << failures << " test(s) FAILED\n"; + return 1; + } + std::cerr << "All plugin helper tests passed.\n"; + return 0; +} diff --git a/tests/e2e/models/smollm3/MODEL.toml b/tests/e2e/models/smollm3/MODEL.toml new file mode 100644 index 0000000000..15e4f1fcc4 --- /dev/null +++ b/tests/e2e/models/smollm3/MODEL.toml @@ -0,0 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +id = "smollm3" +plugin = "smollm3" +test_manifests = [ + "manifests/smollm3-3b.json", +] + +[e2e_defaults.text_generation_causal] +reference_backend = "hf_transformers" +oracle_level = "L1_external_reference" +input_fields = [ + { input = "prompt_repeat", manifest = "prompt_repeat" }, + { input = "expected_prompt_tokens", manifest = "expected_prompt_tokens" }, +] +stages = [ + { name = "full_generation", required = true }, +] diff --git a/tests/e2e/models/smollm3/e2e_plugins/__init__.py b/tests/e2e/models/smollm3/e2e_plugins/__init__.py new file mode 100644 index 0000000000..1f0515e3ff --- /dev/null +++ b/tests/e2e/models/smollm3/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/smollm3/e2e_plugins/comparator.py b/tests/e2e/models/smollm3/e2e_plugins/comparator.py new file mode 100644 index 0000000000..0cf254e8cb --- /dev/null +++ b/tests/e2e/models/smollm3/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 + +"""smollm3 model-owned E2E comparator plugins.""" + +from __future__ import annotations + +from .comparators.text import TextComparator + + +class SmolLM3TextGenerationCausalComparator(TextComparator): + """smollm3 local comparator for text_generation_causal.""" + +comparator = SmolLM3TextGenerationCausalComparator() diff --git a/tests/e2e/models/smollm3/e2e_plugins/comparators/__init__.py b/tests/e2e/models/smollm3/e2e_plugins/comparators/__init__.py new file mode 100644 index 0000000000..9c9857522e --- /dev/null +++ b/tests/e2e/models/smollm3/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/smollm3/e2e_plugins/comparators/_helpers.py b/tests/e2e/models/smollm3/e2e_plugins/comparators/_helpers.py new file mode 100644 index 0000000000..20d4d897b6 --- /dev/null +++ b/tests/e2e/models/smollm3/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/smollm3/e2e_plugins/comparators/text.py b/tests/e2e/models/smollm3/e2e_plugins/comparators/text.py new file mode 100644 index 0000000000..28445ecf3e --- /dev/null +++ b/tests/e2e/models/smollm3/e2e_plugins/comparators/text.py @@ -0,0 +1,511 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Text generation comparator — multi-metric comparison with composite gating. + +Computes logit-level and text-level metrics between TRT and HF reference +outputs and applies composite gating from the ThresholdProfile. No single +metric gates alone; the pass/fail decision uses a composite rule. + +Metrics computed: + 1. logit_cosine_p5 — 5th percentile cosine similarity across steps + 2. logit_rel_l2_p95 — 95th percentile relative L2 norm + 3. stable_top1_match_rate — exact top-1 match where HF margin >= stable_margin + 4. unstable_topk_hit_rate — TRT top-1 in HF top-k where margin < stable_margin + 5. token_agreement_rate — fraction of steps with identical argmax + 6. normalized_text_edit_distance — Levenshtein-normalized on decoded text +""" + +from __future__ import annotations + +import logging +import re +from pathlib import Path + +import numpy as np + +from ..contracts import ( + CompareResult, + MetricResult, + StageOutput, + StageSpec, + StageStatus, + ThresholdProfile, +) +from ._helpers import cosine_similarity, normalized_edit_distance + +logger = logging.getLogger(__name__) + +# Default top-k for unstable token checking +_DEFAULT_TOP_K = 5 +# Default stable margin threshold +_DEFAULT_STABLE_MARGIN = 0.1 + +_COMPOSITE_RULE = ( + "(cosine_p5 >= T OR rel_l2_p95 <= T) " + "AND (agreement >= T OR (stable_top1 >= T AND unstable_topk >= T)) " + "AND ned <= T" +) + +# If the model echoes the prompt, allow a modest amount of non-text preamble +# (warnings/logs) before the prompt appears in stdout. +_PROMPT_SEARCH_MAX_PREFIX_CHARS = 2048 +_MIN_PREFIX_FALLBACK_CHARS = 24 + +# Common multi-turn/chat markers that can appear in decoded text and cause +# cosmetic NED mismatches even when token/logit agreement is strong. +_CHAT_ROLE_PREFIXES = ( + "### response:", + "### assistant:", + "assistant:", + "<|assistant|>", +) +_CHAT_TURN_MARKERS = ( + "### response:", + "### instruction:", + "### assistant:", + "### user:", + "<|assistant|>", + "<|user|>", + "<|im_start|>", + "<|im_end|>", +) + + +def _relative_l2(a: np.ndarray, b: np.ndarray) -> float: + """Relative L2 norm: ||a - b|| / max(||b||, eps).""" + diff_norm = np.linalg.norm(a - b) + ref_norm = np.linalg.norm(b) + return float(diff_norm / max(ref_norm, 1e-12)) + + +def _strip_prompt_echo(text: str, prompt: str) -> str: + """Drop echoed prompt from generated text, tolerating warning preambles. + + Some C++ runs can include tokenizer warnings/log lines before the actual + generated text. If the prompt appears near the beginning of the output, + treat everything before/including it as preamble and compare only the + continuation text. + """ + if not text or not prompt: + return text + + idx = text.find(prompt) + if 0 <= idx <= _PROMPT_SEARCH_MAX_PREFIX_CHARS: + return text[idx + len(prompt):].lstrip() + + return text + + +def _strip_prompt_echo_normalized(text: str, prompt: str) -> str: + """Prompt-echo stripping on normalized text for tokenization-format drift. + + This pass runs after normalization and catches cases where decoded prompt + formatting differs slightly (e.g., whitespace around punctuation), so raw + substring matching misses obvious prompt echoes. + """ + if not text or not prompt: + return text + + norm_prompt = _normalize_for_ned(prompt) + if not norm_prompt: + return text + + if text.startswith(norm_prompt): + return text[len(norm_prompt):].lstrip() + + # Fallback: compare after removing whitespace to handle punctuation-spacing + # drift from tokenizer decode (e.g., "dog. once" vs "dog.once"). + compact_text = "".join(ch for ch in text if not ch.isspace()) + compact_prompt = "".join(ch for ch in norm_prompt if not ch.isspace()) + if compact_prompt and compact_text.startswith(compact_prompt): + remaining = len(compact_prompt) + i = 0 + while i < len(text) and remaining > 0: + if not text[i].isspace(): + remaining -= 1 + i += 1 + return text[i:].lstrip() + + # Keep search window small in normalized space to avoid stripping + # naturally generated prompt repeats that happen later in output. + search_limit = min(_PROMPT_SEARCH_MAX_PREFIX_CHARS, max(256, len(norm_prompt) * 3)) + idx = text.find(norm_prompt) + if 0 <= idx <= search_limit: + return text[idx + len(norm_prompt):].lstrip() + return text + + +def _normalize_for_ned(text: str) -> str: + """Lightweight text normalization before edit-distance comparison.""" + if not text: + return "" + # Collapse whitespace and case-fold to reduce cosmetic diffs. + return " ".join(text.split()).strip().lower() + + +def _strip_leading_role_prefix(text: str) -> str: + """Remove leading chat role prefixes (if present).""" + if not text: + return "" + out = text.lstrip() + while True: + lowered = out.lower() + matched = False + for prefix in _CHAT_ROLE_PREFIXES: + if lowered.startswith(prefix): + out = out[len(prefix):].lstrip() + matched = True + break + if not matched: + return out + + +def _truncate_after_first_turn(text: str) -> str: + """Keep only first assistant turn content and trim trailing markdown stubs.""" + if not text: + return "" + + lowered = text.lower() + cut = len(text) + for marker in _CHAT_TURN_MARKERS: + idx = lowered.find(marker) + if idx > 0: + cut = min(cut, idx) + + out = text[:cut] if cut < len(text) else text + # Some models emit dangling markdown headers (e.g. "##") at the end. + out = re.sub(r"(?:\s*#{2,}\s*)+$", "", out).strip() + return out + + +def _load_logits(stage_output: StageOutput) -> np.ndarray | None: + """Load logits from StageOutput. Returns 2-D array [steps, vocab] or None.""" + # Try logits field first (path or array) + logits = stage_output.logits + if logits is None: + logits = stage_output.data.get("logits_path") + + if logits is None: + return None + + if isinstance(logits, np.ndarray): + return logits + + if isinstance(logits, str) and Path(logits).is_file(): + return np.load(logits) + + return None + + +def _check_numerical_health( + arr: np.ndarray, label: str +) -> list[str]: + """Check for NaN, Inf, and suspicious range. Returns list of warnings.""" + warnings = [] + nan_count = int(np.isnan(arr).sum()) + inf_count = int(np.isinf(arr).sum()) + if nan_count > 0: + warnings.append(f"{label}: {nan_count} NaN values") + if inf_count > 0: + warnings.append(f"{label}: {inf_count} Inf values") + if arr.size > 0: + abs_max = float(np.nanmax(np.abs(arr[np.isfinite(arr)]))) if np.any(np.isfinite(arr)) else 0.0 + if abs_max > 1e6: + warnings.append(f"{label}: large absolute values (max={abs_max:.1e})") + return warnings + + +class TextComparator: + """Multi-metric text generation comparator with composite gating.""" + + @property + def task_strategy(self) -> str: + return "text_generation_causal" + + def compare( + self, + trt: StageOutput, + ref: StageOutput, + threshold: ThresholdProfile, + stage: StageSpec, + ) -> CompareResult: + metrics: dict[str, MetricResult] = {} + + # full_generation runs provide C++ return code from the CLI path. + # If C++ generation failed, surface that explicitly instead of + # allowing debug-runner logits to hide the failure. + cpp_rc = (trt.data or {}).get("cpp_returncode") + if cpp_rc not in (None, 0): + return CompareResult( + stage_name=stage.name, + status=StageStatus.ERROR.value, + metrics=metrics, + message=f"TRT C++ run failed (cpp_returncode={cpp_rc})", + ) + + # Load logits + trt_logits = _load_logits(trt) + ref_logits = _load_logits(ref) + + # Shape/schema check — fall back to text-only for seq2seq models + # where the debug runner doesn't produce logits + if trt_logits is None or ref_logits is None: + return self._compare_text_only(trt, ref, threshold, stage, metrics) + + if trt_logits.ndim != 2 or ref_logits.ndim != 2: + return CompareResult( + stage_name=stage.name, + status=StageStatus.ERROR.value, + metrics=metrics, + message=f"Logits must be 2-D [steps, vocab]: TRT={trt_logits.shape}, HF={ref_logits.shape}", + ) + + # Truncate to common step count + n_steps = min(trt_logits.shape[0], ref_logits.shape[0]) + if n_steps == 0: + return CompareResult( + stage_name=stage.name, + status=StageStatus.ERROR.value, + metrics=metrics, + message="No steps to compare", + ) + + trt_l = trt_logits[:n_steps] + ref_l = ref_logits[:n_steps] + + # Ensure same vocab dimension + notes: list[str] = [] + if trt_l.shape[1] != ref_l.shape[1]: + min_vocab = min(trt_l.shape[1], ref_l.shape[1]) + notes.append( + f"Vocab size mismatch: TRT={trt_l.shape[1]}, HF={ref_l.shape[1]}; " + f"truncating to {min_vocab}" + ) + trt_l = trt_l[:, :min_vocab] + ref_l = ref_l[:, :min_vocab] + + # Numerical health + health_warnings = [] + health_warnings.extend(_check_numerical_health(trt_l, "TRT logits")) + health_warnings.extend(_check_numerical_health(ref_l, "HF logits")) + notes.extend(health_warnings) + + # Replace NaN/Inf with 0 for metric computation + trt_clean = np.nan_to_num(trt_l, nan=0.0, posinf=0.0, neginf=0.0) + ref_clean = np.nan_to_num(ref_l, nan=0.0, posinf=0.0, neginf=0.0) + + thresh = threshold.metrics + + # --- Metric 1: logit_cosine_p5 --- + cosines = np.array([ + cosine_similarity(trt_clean[i], ref_clean[i]) + for i in range(n_steps) + ]) + logit_cosine_p5 = float(np.percentile(cosines, 5)) + cosine_thresh = thresh.get("logit_cosine_p5", 0.99) + metrics["logit_cosine_p5"] = MetricResult( + value=logit_cosine_p5, threshold=cosine_thresh, + operator=">=", passed=logit_cosine_p5 >= cosine_thresh, + ) + + # --- Metric 2: logit_rel_l2_p95 --- + rel_l2s = np.array([ + _relative_l2(trt_clean[i], ref_clean[i]) + for i in range(n_steps) + ]) + logit_rel_l2_p95 = float(np.percentile(rel_l2s, 95)) + rel_l2_thresh = thresh.get("logit_rel_l2_p95", 0.05) + metrics["logit_rel_l2_p95"] = MetricResult( + value=logit_rel_l2_p95, threshold=rel_l2_thresh, + operator="<=", passed=logit_rel_l2_p95 <= rel_l2_thresh, + ) + + # --- Per-step argmax and margin analysis --- + trt_argmax = trt_clean.argmax(axis=1) + ref_argmax = ref_clean.argmax(axis=1) + + ref_sorted = np.sort(ref_clean, axis=1) + hf_margin = ref_sorted[:, -1] - ref_sorted[:, -2] + + stable_margin = thresh.get("stable_margin", _DEFAULT_STABLE_MARGIN) + top_k = int(thresh.get("top_k", _DEFAULT_TOP_K)) + + stable_mask = hf_margin >= stable_margin + unstable_mask = ~stable_mask + n_stable = int(stable_mask.sum()) + n_unstable = int(unstable_mask.sum()) + + # --- Metric 3: stable_top1_match_rate --- + if n_stable > 0: + stable_matches = int((trt_argmax[stable_mask] == ref_argmax[stable_mask]).sum()) + stable_top1_match_rate = stable_matches / n_stable + else: + stable_top1_match_rate = 1.0 + stable_thresh = thresh.get("stable_top1_match_rate", 0.9) + metrics["stable_top1_match_rate"] = MetricResult( + value=stable_top1_match_rate, threshold=stable_thresh, + operator=">=", passed=stable_top1_match_rate >= stable_thresh, + note=f"{n_stable} stable steps", + ) + + # --- Metric 4: unstable_topk_hit_rate --- + if n_unstable > 0: + ref_topk = np.argsort(ref_clean, axis=1)[:, -top_k:] + hits = 0 + unstable_indices = np.where(unstable_mask)[0] + for idx in unstable_indices: + if trt_argmax[idx] in ref_topk[idx]: + hits += 1 + unstable_topk_hit_rate = hits / n_unstable + else: + unstable_topk_hit_rate = 1.0 + topk_thresh = thresh.get("unstable_topk_hit_rate", 0.8) + metrics["unstable_topk_hit_rate"] = MetricResult( + value=unstable_topk_hit_rate, threshold=topk_thresh, + operator=">=", passed=unstable_topk_hit_rate >= topk_thresh, + note=f"{n_unstable} unstable steps", + ) + + # --- Metric 5: token_agreement_rate --- + token_agreement_rate = float((trt_argmax == ref_argmax).mean()) + ta_thresh = thresh.get("token_agreement_rate", 0.8) + metrics["token_agreement_rate"] = MetricResult( + value=token_agreement_rate, threshold=ta_thresh, + operator=">=", passed=token_agreement_rate >= ta_thresh, + ) + + # --- Metric 6: normalized_text_edit_distance --- + trt_text = (trt.text or "").strip() + ref_text = (ref.text or "").strip() + + prompt = (trt.data or {}).get("prompt", "") + # Prompt echo handling is TRT-side only. HF reference text is decoded + # from generated tokens and should not include prompt prefill; stripping + # prompt from reference can incorrectly remove legitimate generated text + # if the model naturally repeats the prompt phrase later. + trt_text_for_ned = _normalize_for_ned( + _truncate_after_first_turn( + _strip_leading_role_prefix(_strip_prompt_echo(trt_text, prompt)) + ) + ) + ref_text_for_ned = _normalize_for_ned( + _truncate_after_first_turn( + _strip_leading_role_prefix(ref_text) + ) + ) + trt_text_for_ned = _strip_prompt_echo_normalized(trt_text_for_ned, prompt) + # Seq2seq models output text that may start with the prompt (e.g. + # BART reconstructing its input). Strip the prompt prefix from ref + # only when it appears at the very start of the normalized text. + # This avoids accidentally removing prompt substrings that appear + # later in naturally generated text from causal models. + norm_prompt = _normalize_for_ned(prompt) + if norm_prompt and ref_text_for_ned.startswith(norm_prompt): + ref_text_for_ned = ref_text_for_ned[len(norm_prompt):].lstrip() + + if trt_text_for_ned or ref_text_for_ned: + ned = normalized_edit_distance(trt_text_for_ned, ref_text_for_ned) + # Some TRT CLI paths stop decoding early on EOS while the debug/HF + # text path keeps fixed-length continuation tokens. If token/logit + # metrics already agree, compare on the common prefix to avoid + # false NED hard-fails caused purely by suffix length mismatch. + ta_thresh = thresh.get("token_agreement_rate", 0.8) + if token_agreement_rate >= ta_thresh: + if len(trt_text_for_ned) <= len(ref_text_for_ned): + short, long = trt_text_for_ned, ref_text_for_ned + else: + short, long = ref_text_for_ned, trt_text_for_ned + if len(short) >= _MIN_PREFIX_FALLBACK_CHARS and long.startswith(short): + prefix_ned = normalized_edit_distance(short, long[:len(short)]) + if prefix_ned < ned: + notes.append( + "NED prefix fallback applied (matching continuation prefix; " + "suffix length mismatch likely due EOS stopping behavior)" + ) + ned = prefix_ned + else: + ned = 0.0 + ned_thresh = thresh.get("normalized_text_edit_distance", 0.2) + metrics["normalized_text_edit_distance"] = MetricResult( + value=ned, threshold=ned_thresh, + operator="<=", passed=ned <= ned_thresh, + ) + + # --- Composite gating --- + logit_quality_ok = ( + metrics["logit_cosine_p5"].passed + or metrics["logit_rel_l2_p95"].passed + ) + + token_level_ok = ( + metrics["token_agreement_rate"].passed + or ( + metrics["stable_top1_match_rate"].passed + and metrics["unstable_topk_hit_rate"].passed + ) + ) + + text_ok = metrics["normalized_text_edit_distance"].passed + + passed = logit_quality_ok and token_level_ok and text_ok + + message = ( + f"{'PASS' if passed else 'FAIL'}: " + f"cosine_p5={logit_cosine_p5:.4f}, " + f"agreement={token_agreement_rate:.4f}, " + f"ned={ned:.4f}" + ) + + return CompareResult( + stage_name=stage.name, + status=StageStatus.PASSED.value if passed else StageStatus.FAILED.value, + metrics=metrics, + composite_rule=_COMPOSITE_RULE, + message=message, + ) + + + def _compare_text_only( + self, + trt: StageOutput, + ref: StageOutput, + threshold: ThresholdProfile, + stage: StageSpec, + metrics: dict[str, MetricResult], + ) -> CompareResult: + """Text-only comparison when logits are unavailable (seq2seq models).""" + thresh = threshold.metrics + + prompt = (trt.data or {}).get("prompt", "") + trt_text = _normalize_for_ned( + _strip_leading_role_prefix(_strip_prompt_echo((trt.text or "").strip(), prompt)) + ) + ref_text = _normalize_for_ned( + _strip_leading_role_prefix(_strip_prompt_echo((ref.text or "").strip(), prompt)) + ) + + if trt_text and ref_text: + ned = normalized_edit_distance(trt_text, ref_text) + elif not trt_text and not ref_text: + ned = 0.0 + else: + ned = 1.0 + + ned_thresh = thresh.get("normalized_text_edit_distance", 0.2) + metrics["normalized_text_edit_distance"] = MetricResult( + value=ned, threshold=ned_thresh, + operator="<=", passed=ned <= ned_thresh, + ) + + passed = ned <= ned_thresh + return CompareResult( + stage_name=stage.name, + status=StageStatus.PASSED.value if passed else StageStatus.FAILED.value, + metrics=metrics, + composite_rule="text-only (logits unavailable for seq2seq): ned <= threshold", + message=f"{'PASS' if passed else 'FAIL'}: text-only ned={ned:.4f}", + ) + + +plugin = TextComparator() diff --git a/tests/e2e/models/smollm3/e2e_plugins/contract.py b/tests/e2e/models/smollm3/e2e_plugins/contract.py new file mode 100644 index 0000000000..fdfac7d998 --- /dev/null +++ b/tests/e2e/models/smollm3/e2e_plugins/contract.py @@ -0,0 +1,262 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""smollm3-owned E2E contract plugins.""" +from __future__ import annotations + +import re + +from tests.e2e_harness.contracts import ( + CompareResult, + MetricResult, + StageStatus, +) + + +_PREFILL_OBSERVATION_RE = re.compile( + r"^\[trtmc\.prefill\] tokens=(\d+) launches=(\d+) max_chunk=(\d+)$" +) + + +def contract_config(case): + config = case.metadata.get("contract_config", {}) + return dict(config) if isinstance(config, dict) else {} + + +def normalize_text(text: str) -> str: + if not text: + return "" + return " ".join(text.split()).strip().lower() + + +def strip_prompt_echo(text: str, prompt: str) -> str: + if not text or not prompt: + return text + idx = text.find(prompt) + if 0 <= idx <= 2048: + return text[idx + len(prompt):].lstrip() + norm_text = normalize_text(text) + norm_prompt = normalize_text(prompt) + if norm_prompt and norm_text.startswith(norm_prompt): + return text[len(prompt):].lstrip() if text.startswith(prompt) else text + return text + + +_CHAT_ROLE_PREFIXES = ( + "### response:", "### assistant:", "assistant:", + "<|assistant|>", "<|im_start|>assistant\n", +) + +_CHAT_TURN_MARKERS = ( + "### response:", "### instruction:", "### assistant:", + "### user:", "<|assistant|>", "<|user|>", + "<|im_start|>", "<|im_end|>", +) + + +def strip_chat_markup(text: str) -> str: + if not text: + return "" + out = text.lstrip() + while True: + lowered = out.lower() + matched = False + for prefix in _CHAT_ROLE_PREFIXES: + if lowered.startswith(prefix): + out = out[len(prefix):].lstrip() + matched = True + break + if not matched: + break + lowered = out.lower() + cut = len(out) + for marker in _CHAT_TURN_MARKERS: + idx = lowered.find(marker) + if idx > 0: + cut = min(cut, idx) + if cut < len(out): + out = out[:cut] + out = re.sub(r"(?:\s*#{2,}\s*)+$", "", out).strip() + return out + + +def extract_answer(output, prompt: str = "") -> str: + raw = output.text or "" + if prompt: + raw = strip_prompt_echo(raw, prompt) + raw = strip_chat_markup(raw) + return raw.strip() + + +def levenshtein_ned(a: str, b: str) -> float: + if not a and not b: + return 0.0 + max_len = max(len(a), len(b)) + if max_len == 0: + return 0.0 + if len(a) < len(b): + a, b = b, a + prev = list(range(len(b) + 1)) + for i, c1 in enumerate(a): + curr = [i + 1] + for j, c2 in enumerate(b): + curr.append(min( + prev[j + 1] + 1, + curr[j] + 1, + prev[j] + (0 if c1 == c2 else 1), + )) + prev = curr + return prev[-1] / max_len + + +def make_pass(stage_name: str, metrics, rule: str = ""): + from tests.e2e_harness.contracts import CompareResult + return CompareResult( + stage_name=stage_name, + status="passed", + metrics=metrics, + composite_rule=rule, + message="Contract verified", + ) + + +def make_fail(stage_name: str, metrics, rule: str = "", message: str = ""): + from tests.e2e_harness.contracts import CompareResult + return CompareResult( + stage_name=stage_name, + status="failed", + metrics=metrics, + composite_rule=rule, + message=message or "Contract verification failed", + ) + + +def make_skip(stage_name: str, metrics, rule: str = "", message: str = ""): + from tests.e2e_harness.contracts import CompareResult + return CompareResult( + stage_name=stage_name, + status="skipped", + metrics=metrics, + composite_rule=rule, + message=message or "Contract validation skipped", + ) + + +def make_error(stage_name: str, error: str): + from tests.e2e_harness.contracts import CompareResult + return CompareResult( + stage_name=stage_name, + status="error", + message=f"Contract verification error: {error}", + ) + +class SmolLM3CausalContinuationPlugin: + reference_families = ["causal_base_continuation"] + user_contract = "continuation_parity" + + def configure_reference(self, case): + return contract_config(case) + + def verify(self, trt_output, ref_output, case, threshold): + cpp_rc = (trt_output.data or {}).get("cpp_returncode") + if cpp_rc not in (None, 0): + metrics = { + "cpp_returncode_ok": MetricResult( + value=0.0, + threshold=1.0, + operator="==", + passed=False, + note=f"cpp_returncode={cpp_rc}", + ), + } + detail = (trt_output.data or {}).get("cpp_runtime_error") + suffix = f": {detail}" if detail else "" + return CompareResult( + stage_name="full_generation", + status=StageStatus.ERROR.value, + metrics=metrics, + message=f"TRT C++ run failed (cpp_returncode={cpp_rc}){suffix}", + ) + + prompt = case.inputs.get("prompt", "") + config = contract_config(case) + preserve_prompt_echo = bool(config.get("preserve_prompt_echo")) + reconstruction_check = bool(config.get("seq2seq_reconstruction")) + if preserve_prompt_echo: + trt_text = normalize_text(trt_output.text or "") + ref_text = normalize_text(ref_output.text or "") + else: + trt_text = normalize_text(strip_prompt_echo(trt_output.text or "", prompt)) + ref_text = normalize_text(strip_prompt_echo(ref_output.text or "", prompt)) + + if not trt_text and not ref_text: + metrics = { + "non_empty_continuation": MetricResult( + value=0.0, + threshold=1.0, + operator="==", + passed=False, + note="empty TRT and reference text do not validate parity", + ), + } + return make_fail( + "full_generation", + metrics, + "non-empty continuation required", + "Both TRT and reference produced empty continuation", + ) + + ned = levenshtein_ned(trt_text, ref_text) + ned_threshold = threshold.metrics.get("contract_ned_threshold", 0.25) + prefix_len = min(50, min(len(trt_text), len(ref_text))) + prefix_match = (trt_text[:prefix_len] == ref_text[:prefix_len]) if prefix_len > 0 else True + + metrics = { + "ned": MetricResult( + value=ned, + threshold=ned_threshold, + operator="<=", + passed=ned <= ned_threshold, + ), + "prefix_match": MetricResult( + value=1.0 if prefix_match else 0.0, + threshold=1.0, + operator="==", + passed=prefix_match, + note=f"first {prefix_len} chars", + ), + } + if reconstruction_check: + metrics["non_empty_trt_text"] = MetricResult( + value=1.0 if trt_text else 0.0, + threshold=1.0, + operator="==", + passed=bool(trt_text), + note="visible TRT reconstruction text", + ) + metrics["non_empty_reference_text"] = MetricResult( + value=1.0 if ref_text else 0.0, + threshold=1.0, + operator="==", + passed=bool(ref_text), + note="visible HF reconstruction text", + ) + + passed = ned <= ned_threshold + rule = ( + "seq2seq reconstruction parity against HF reference" + if reconstruction_check + else "ned <= threshold (continuation parity)" + ) + if passed: + return make_pass("full_generation", metrics, rule) + return make_fail( + "full_generation", + metrics, + rule, + f"Continuation diverged: NED={ned:.3f}", + ) + +plugin = [ + SmolLM3CausalContinuationPlugin(), +] diff --git a/tests/e2e/models/smollm3/e2e_plugins/contracts.py b/tests/e2e/models/smollm3/e2e_plugins/contracts.py new file mode 100644 index 0000000000..d6f9281d2a --- /dev/null +++ b/tests/e2e/models/smollm3/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/smollm3/e2e_plugins/reference.py b/tests/e2e/models/smollm3/e2e_plugins/reference.py new file mode 100644 index 0000000000..746d0fa5e1 --- /dev/null +++ b/tests/e2e/models/smollm3/e2e_plugins/reference.py @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""smollm3 model-owned E2E reference plugins.""" + +from __future__ import annotations + +from .references.hf_transformers import HfTransformersReference + + +class SmolLM3HfTransformersReference(HfTransformersReference): + """smollm3 local reference for hf_transformers.""" + + +reference = [SmolLM3HfTransformersReference()] diff --git a/tests/e2e/models/smollm3/e2e_plugins/references/__init__.py b/tests/e2e/models/smollm3/e2e_plugins/references/__init__.py new file mode 100644 index 0000000000..6487ae8000 --- /dev/null +++ b/tests/e2e/models/smollm3/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/smollm3/e2e_plugins/references/hf_transformers.py b/tests/e2e/models/smollm3/e2e_plugins/references/hf_transformers.py new file mode 100644 index 0000000000..0c455e45f2 --- /dev/null +++ b/tests/e2e/models/smollm3/e2e_plugins/references/hf_transformers.py @@ -0,0 +1,984 @@ +# 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) + 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_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/smollm3/e2e_plugins/runner.py b/tests/e2e/models/smollm3/e2e_plugins/runner.py new file mode 100644 index 0000000000..c1590d9cfe --- /dev/null +++ b/tests/e2e/models/smollm3/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 + +"""smollm3 model-owned E2E runner plugins.""" + +from __future__ import annotations + +from .runners.text_generation import TextGenerationCausalRunner + + +class SmolLM3TextGenerationCausalRunner(TextGenerationCausalRunner): + """smollm3 local runner for text_generation_causal.""" + +runner = SmolLM3TextGenerationCausalRunner() diff --git a/tests/e2e/models/smollm3/e2e_plugins/runners/__init__.py b/tests/e2e/models/smollm3/e2e_plugins/runners/__init__.py new file mode 100644 index 0000000000..986743e50e --- /dev/null +++ b/tests/e2e/models/smollm3/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/smollm3/e2e_plugins/runners/text_generation.py b/tests/e2e/models/smollm3/e2e_plugins/runners/text_generation.py new file mode 100644 index 0000000000..286584ef10 --- /dev/null +++ b/tests/e2e/models/smollm3/e2e_plugins/runners/text_generation.py @@ -0,0 +1,929 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Text generation causal strategy runner -- TRT inference via C++ binary and debug runner. + +Handles smollm3's smollm3_decoder_kv_cache runtime strategy, which maps to +task_strategy="text_generation_causal". + +Supported stages: + - "full_generation": C++ binary inference + debug runner logits (both prefill + decode) + - "prefill": Debug runner prefill-only (per input-token logits) + - "decode": Debug runner decode-only (per generated-token logits, assumes prefill done) + +All GPU work runs in subprocesses to prevent OOM when testing multiple models. +""" + +from __future__ import annotations + +import logging +import json +import os +import re +import shutil +import subprocess +import sys +import tempfile +import textwrap +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__) + +_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 _prompt_from_case(case: E2ECase, ctx: RunContext | None = None) -> str: + repeat = case.inputs.get("prompt_repeat") + if isinstance(repeat, dict): + text = str(repeat["text"]) + separator = str(repeat.get("separator", "")) + count = int(repeat["count"]) + prompt = separator.join([text] * count) + str(repeat.get("suffix", "")) + if ctx is not None and ctx.artifacts_dir: + output_dir = Path(_case_artifact_dir(ctx.artifacts_dir, case.name)) + output_dir.mkdir(parents=True, exist_ok=True) + (output_dir / "resolved_prompt.txt").write_text(prompt, encoding="utf-8") + return prompt + return str(case.inputs.get("prompt", "The capital of France is")) + + +def _count_prompt_tokens( + case: E2ECase, + ctx: RunContext, + prompt: str, +) -> tuple[int, dict]: + python = ctx.runtime_python_path() or sys.executable + script = textwrap.dedent( + """\ + import sys + from transformers import AutoTokenizer + + revision = sys.argv[2] or None + tokenizer = AutoTokenizer.from_pretrained( + sys.argv[1], + revision=revision, + trust_remote_code=sys.argv[3] == "1", + ) + print(len(tokenizer.encode(sys.stdin.read(), add_special_tokens=False))) + """ + ) + cmd = [ + python, + "-c", + script, + case.hf_id, + case.hf_revision, + "1" if case.metadata.get("trust_remote_code") else "0", + ] + result = subprocess.run( + cmd, + input=prompt, + capture_output=True, + text=True, + timeout=120, + ) + meta = { + "command": cmd, + "returncode": result.returncode, + "stdout": result.stdout, + "stderr": result.stderr, + } + if result.returncode != 0: + raise RuntimeError( + "Failed to count prompt tokens with the pinned model tokenizer: " + f"{result.stderr.strip()}" + ) + try: + return int(result.stdout.strip()), meta + except ValueError as exc: + raise RuntimeError( + f"Tokenizer returned an invalid prompt token count: {result.stdout!r}" + ) from exc + + +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}" + ) + + +class TextGenerationCausalRunner: + """Execute TRT text generation inference via C++ binary + Python debug runner.""" + + @property + def strategy_name(self) -> str: + return "text_generation_causal" + + 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 == "prefill": + return self._run_prefill(case, stage, ctx) + if stage.name == "decode": + return self._run_decode(case, stage, ctx) + raise ValueError( + f"Unknown stage {stage.name!r} for text_generation_causal. " + f"Supported: {_SUPPORTED_STAGES}" + ) + + # ------------------------------------------------------------------ + # full_generation: C++ binary + debug runner (prefill + decode) + # ------------------------------------------------------------------ + + def _run_full_generation( + self, case: E2ECase, stage: StageSpec, ctx: RunContext + ) -> StageOutput: + """Run C++ binary inference and capture per-step logits via debug runner.""" + bundle_path = str(Path(ctx.engine_dir) / case.bundle) + prompt = _prompt_from_case(case, ctx) + max_new_tokens = case.inputs.get("max_new_tokens", 30) + prompt_token_count = None + prompt_validation_meta = None + if "expected_prompt_tokens" in case.inputs: + prompt_token_count, prompt_validation_meta = _count_prompt_tokens( + case, ctx, prompt + ) + expected_prompt_tokens = int(case.inputs["expected_prompt_tokens"]) + if prompt_token_count != expected_prompt_tokens: + raise RuntimeError( + f"Prompt fixture token count changed: expected " + f"{expected_prompt_tokens}, got {prompt_token_count}" + ) + has_contract = bool(case.reference_family and case.user_contract) + is_acceptance = case.ci_lane == "acceptance" + + use_single_process_debug = bool( + case.metadata.get("single_process_debug_generation", False) + ) and not (has_contract and is_acceptance) + if use_single_process_debug: + logits_path, debug_time, debug_meta = self._run_debug_runner_logits( + ctx, bundle_path, prompt, max_new_tokens, case, phase="full" + ) + text = str(debug_meta.get("full_text") or debug_meta.get("generated_text") or "") + cpp_rc = int(debug_meta.get("returncode", -1)) + data = { + "cpp_text": text, + "cpp_returncode": cpp_rc, + "prompt": prompt, + "runner_mode": "single_process_debug_generation", + } + if prompt_token_count is not None: + data["prompt_token_count"] = prompt_token_count + if logits_path: + data["logits_path"] = logits_path + return StageOutput( + stage_name=stage.name, + data=data, + text=text, + logits=logits_path, + timing_s=debug_time, + metadata={ + "cpp": {"skipped": "single_process_debug_generation"}, + "debug_runner": debug_meta, + "prompt_validation": prompt_validation_meta, + }, + ) + + # C++ binary inference + cpp_text, cpp_time, cpp_meta = self._run_cpp_binary( + ctx, bundle_path, prompt, max_new_tokens, case=case, inputs=case.inputs + ) + + # Debug runner for per-step logits — skip in acceptance lane when + # a contract plugin handles verification (only needs text, not logits) + skip_debug = has_contract and is_acceptance + + if skip_debug: + logits_path = None + debug_time = 0.0 + debug_meta = {"skipped": "contract plugin active in acceptance lane"} + else: + logits_path, debug_time, debug_meta = self._run_debug_runner_logits( + ctx, bundle_path, prompt, max_new_tokens, case, phase="full" + ) + if logits_path is None and _distributed_debug_logits_required(case): + raise RuntimeError(_format_debug_runner_error(case, "full", debug_meta)) + + data = { + "cpp_text": cpp_text, + "cpp_returncode": cpp_meta.get("effective_returncode", cpp_meta.get("returncode", -1)), + "prompt": prompt, + } + if prompt_token_count is not None: + data["prompt_token_count"] = prompt_token_count + if cpp_meta.get("runtime_error_detected"): + data["cpp_runtime_error"] = cpp_meta["runtime_error_detected"] + if cpp_meta.get("token_ids") is not None: + data["token_ids"] = cpp_meta["token_ids"] + if cpp_meta.get("text_output_path"): + data["text_output_path"] = cpp_meta["text_output_path"] + contract_config = case.metadata.get("contract_config", {}) + if "token_parity_ignore_terminal_token_ids" in contract_config: + data["token_parity_ignore_terminal_token_ids"] = ( + contract_config["token_parity_ignore_terminal_token_ids"] + ) + if "token_parity_eos_token_ids" in contract_config: + data["token_parity_eos_token_ids"] = contract_config["token_parity_eos_token_ids"] + if "forbidden_token_ids" in contract_config: + data["forbidden_token_ids"] = contract_config["forbidden_token_ids"] + if logits_path: + data["logits_path"] = logits_path + + return StageOutput( + stage_name=stage.name, + data=data, + text=cpp_text, + logits=logits_path, + timing_s=cpp_time + debug_time, + metadata={ + "cpp": cpp_meta, + "debug_runner": debug_meta, + "prompt_validation": prompt_validation_meta, + }, + ) + + # ------------------------------------------------------------------ + # prefill: debug runner prefill-only (per input-token logits) + # ------------------------------------------------------------------ + + def _run_prefill( + self, case: E2ECase, stage: StageSpec, ctx: RunContext + ) -> StageOutput: + """Run debug runner prefill phase only -- logits for each input token.""" + bundle_path = str(Path(ctx.engine_dir) / case.bundle) + prompt = _prompt_from_case(case, ctx) + + logits_path, elapsed, meta = self._run_debug_runner_logits( + ctx, bundle_path, prompt, max_new_tokens=0, case=case, phase="prefill" + ) + if logits_path is None and _distributed_debug_logits_required(case): + raise RuntimeError(_format_debug_runner_error(case, "prefill", meta)) + + data = {} + if logits_path: + data["logits_path"] = logits_path + + return StageOutput( + stage_name=stage.name, + data=data, + text=None, + logits=logits_path, + timing_s=elapsed, + metadata={"debug_runner": meta}, + ) + + # ------------------------------------------------------------------ + # decode: debug runner decode-only (per generated-token logits) + # ------------------------------------------------------------------ + + def _run_decode( + self, case: E2ECase, stage: StageSpec, ctx: RunContext + ) -> StageOutput: + """Run debug runner decode phase only -- logits for generated tokens.""" + bundle_path = str(Path(ctx.engine_dir) / case.bundle) + prompt = _prompt_from_case(case, ctx) + max_new_tokens = case.inputs.get("max_new_tokens", 30) + + logits_path, elapsed, meta = self._run_debug_runner_logits( + ctx, bundle_path, prompt, max_new_tokens, case=case, phase="decode" + ) + if logits_path is None and _distributed_debug_logits_required(case): + raise RuntimeError(_format_debug_runner_error(case, "decode", meta)) + + data = {} + if logits_path: + data["logits_path"] = logits_path + + return StageOutput( + stage_name=stage.name, + data=data, + text=None, + logits=logits_path, + timing_s=elapsed, + metadata={"debug_runner": meta}, + ) + + # ------------------------------------------------------------------ + # Subprocess helpers + # ------------------------------------------------------------------ + + def _run_cpp_binary( + self, + ctx: RunContext, + bundle_path: str, + prompt: str, + max_new_tokens: int, + case: E2ECase | None = None, + inputs: dict | None = None, + ) -> tuple[str, float, dict]: + """Run the C++ trtmc binary as a subprocess. Returns (text, time_s, meta).""" + cmd = [ + ctx.binary_path, "run", bundle_path, + "--prompt", prompt, + "--max-new-tokens", str(max_new_tokens), + ] + output_jsonl_path: Path | None = None + if case is not None and not _distributed_runtime_config(case): + output_root = ( + Path(_case_artifact_dir(ctx.artifacts_dir, case.name)) + if ctx.artifacts_dir + else Path(tempfile.gettempdir()) + ) + output_root.mkdir(parents=True, exist_ok=True) + output_jsonl_path = output_root / "trt_text_generation.jsonl" + cmd.extend(["-o", str(output_jsonl_path)]) + runtime_cli_python = ctx.runtime_cli_hf_python() + if runtime_cli_python: + cmd.extend(["--hf-python", runtime_cli_python]) + if inputs: + if inputs.get("temperature", 1.0) != 1.0: + cmd.extend(["--temperature", str(inputs["temperature"])]) + if inputs.get("top_p", 1.0) < 1.0 - 1e-6: + cmd.extend(["--top-p", str(inputs["top_p"])]) + if inputs.get("min_p", 0.0) > 1e-6: + cmd.extend(["--min-p", str(inputs["min_p"])]) + if inputs.get("top_k", 1) != 1: + cmd.extend(["--top-k", str(inputs["top_k"])]) + if inputs.get("seed", -1) >= 0: + cmd.extend(["--seed", str(inputs["seed"])]) + if inputs.get("generation_mode"): + cmd.extend(["--generation-mode", str(inputs["generation_mode"])]) + if inputs.get("block_length", 0): + cmd.extend(["--block-length", str(inputs["block_length"])]) + if inputs.get("threshold") is not None: + cmd.extend(["--threshold", str(inputs["threshold"])]) + + if case is not None: + contract_config = case.metadata.get("contract_config", {}) + if contract_config.get("use_chat_template"): + cmd.append("--chat-template") + if contract_config.get("enable_thinking") is False: + cmd.append("--no-thinking") + + env = dict(os.environ) + if ctx.ld_library_path: + env["LD_LIBRARY_PATH"] = ctx.ld_library_path + distributed_runtime = _distributed_runtime_config(case) + if distributed_runtime and case is not None: + _ensure_distributed_runtime_env(case, ctx, env) + extra_env = distributed_runtime.get("env", {}) + if isinstance(extra_env, dict): + env.update({str(k): str(v) for k, v in extra_env.items()}) + cmd = _wrap_distributed_command(cmd, case, env) + + logger.info("C++ inference: %s", " ".join(cmd)) + t0 = time.monotonic() + memory_sampler = _maybe_start_gpu_memory_sampler(distributed_runtime, ctx, case, env) + try: + result = subprocess.run( + cmd, capture_output=True, text=True, timeout=600, env=env + ) + except subprocess.TimeoutExpired: + elapsed = time.monotonic() - t0 + meta = {"returncode": -1, "error": "timeout"} + if memory_sampler is not None: + meta["gpu_memory"] = memory_sampler.stop() + return "", elapsed, meta + except Exception as e: + elapsed = time.monotonic() - t0 + meta = {"returncode": -1, "error": str(e)} + if memory_sampler is not None: + meta["gpu_memory"] = memory_sampler.stop() + return "", elapsed, meta + elapsed = time.monotonic() - t0 + memory_meta = memory_sampler.stop() if memory_sampler is not None else None + + parse_stderr = _strip_mpi_stream_tags(result.stderr) if distributed_runtime else result.stderr + meta: dict = { + "returncode": result.returncode, + "command": cmd, + "stdout": result.stdout, + "stderr": result.stderr, + } + if distributed_runtime: + meta["distributed_runtime"] = distributed_runtime + meta["rank_zero_stdout"] = _extract_rank_zero_stdout(result.stdout) + meta["stderr_without_mpi_tags"] = parse_stderr + if memory_meta is not None: + meta["gpu_memory"] = memory_meta + meta.update(_extract_trtmc_timing(parse_stderr)) + meta.update(_extract_trtmc_load_timing(parse_stderr)) + runtime_error = _detect_trt_runtime_error(parse_stderr) + if runtime_error: + meta["runtime_error_detected"] = runtime_error + if result.returncode == 0: + meta["effective_returncode"] = -1 + meta["error"] = "TensorRT runtime error detected in stderr" + + if result.returncode != 0 or runtime_error: + truncated, log_path = save_full_stderr( + result.stderr, ctx.artifacts_dir or "", "cpp_binary") + meta["stderr_truncated"] = truncated + if log_path: + meta["stderr_log"] = log_path + + text = _extract_rank_zero_stdout(result.stdout) if distributed_runtime else result.stdout.strip() + if output_jsonl_path is not None: + sample = _read_text_generation_sample(output_jsonl_path) + if sample: + meta["text_output_path"] = str(output_jsonl_path) + if isinstance(sample.get("generated"), str): + text = sample["generated"] + meta["generated"] = text + if isinstance(sample.get("token_ids"), list): + meta["token_ids"] = sample["token_ids"] + return text, elapsed, meta + + def _run_debug_runner_logits( + self, + ctx: RunContext, + bundle_path: str, + prompt: str, + max_new_tokens: int, + case: E2ECase, + phase: str = "full", + ) -> tuple[str | None, float, dict]: + """Run TrtRunner in a subprocess to collect per-step logits. + + Args: + phase: "full" = prefill + decode, "prefill" = input tokens only, + "decode" = generated tokens only (still runs prefill internally). + + Returns (logits_npy_path, time_s, meta). + """ + 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) / f"trt_{phase}_logits.npy" + ) + + script = textwrap.dedent(f"""\ + import sys, json, numpy as np + from pathlib import Path + + bundle_path = {bundle_path!r} + prompt = {prompt!r} + max_new_tokens = {max_new_tokens} + logits_path = {logits_path!r} + phase = {phase!r} + distributed = {bool(_distributed_runtime_config(case))!r} + tp_size = {int(_distributed_runtime_config(case).get("world_size", _distributed_runtime_config(case).get("tp_size", 1)) or 1)} + + # Create the family-owned runner from bundle metadata. + from tensorrt_model_connect.debug_runner import ( + TensorParallelNcclGroup, + ) + from tensorrt_model_connect.families.smollm3.debug_runner import ( + load_config_from_bundle, + load_engine_from_bundle, + runner_from_bundle as family_runner_from_bundle, + ) + from tensorrt_model_connect.parallel_config import rank_engine_section + group = None + runner = None + try: + config_json = load_config_from_bundle(bundle_path) + engine_section = "engine_plan" + distributed_communicator = None + if distributed: + group = TensorParallelNcclGroup(world_size=tp_size) + engine_section = rank_engine_section(group.rank) + distributed_communicator = group.communicator + engine_plan, header = load_engine_from_bundle( + bundle_path, section_name=engine_section) + runner = family_runner_from_bundle( + runtime_strategy=str(config_json.get("runtime_strategy") or ""), + config=config_json, + header=header, + engine_plan=engine_plan, + bundle_path=bundle_path, + distributed_communicator=distributed_communicator, + ) + + # Tokenize + from transformers import AutoTokenizer + hf_id = config_json.get("_hf_id", {case.hf_id!r}) + trust_remote_code = {case.metadata.get("trust_remote_code", False)!r} + tokenizer = AutoTokenizer.from_pretrained( + hf_id, trust_remote_code=trust_remote_code) + input_ids = tokenizer.encode(prompt) + + # Run full generate (we always need prefill internally) + results = runner.generate(input_ids, max_new_tokens) + is_seq2seq = runner.__class__.__name__ == "Seq2SeqTrtRunner" + generated_tokens = [] + if len(results) > 0 and max_new_tokens > 0: + start = 0 if is_seq2seq else max(len(input_ids) - 1, 0) + for i in range(max_new_tokens): + idx = start + i + if idx >= len(results): + break + generated_tokens.append( + int(np.argmax(results[idx]["logits"].flatten())) + ) + full_ids = input_ids + generated_tokens + generated_text = tokenizer.decode( + generated_tokens, skip_special_tokens=True) + full_text = tokenizer.decode(full_ids, skip_special_tokens=True) + + # Select phase slice + n_input = len(input_ids) + if phase == "prefill": + results = results[:n_input] + elif phase == "decode": + results = results[n_input:] + # else "full": keep all + + logits_list = [r["logits"].flatten() for r in results] + + should_write = group is None or group.rank == 0 + rank = 0 if group is None else group.rank + if len(logits_list) == 0: + if should_write: + np.save(logits_path, np.zeros((0, 0), dtype=np.float32)) + print(f"OK rank={{rank}} steps=0 vocab=0") + else: + max_len = max(l.shape[0] for l in logits_list) + padded = np.zeros((len(logits_list), max_len), dtype=np.float32) + for i, l in enumerate(logits_list): + padded[i, :l.shape[0]] = l + if should_write: + np.save(logits_path, padded) + print(f"OK rank={{rank}} steps={{len(logits_list)}} vocab={{max_len}}") + if should_write: + print("TRTMC_DEBUG_META " + json.dumps({{ + "generated_text": generated_text, + "full_text": full_text, + "generated_token_count": len(generated_tokens), + "distributed_rank": rank, + }})) + finally: + if runner is not None: + del runner + runner = None + if group is not None: + group.close() + """) + + python = ctx.runtime_python_path() or sys.executable + logger.info("Debug runner (%s): collecting logits for %s", phase, case.name) + env = dict(os.environ) + if ctx.ld_library_path: + env["LD_LIBRARY_PATH"] = ctx.ld_library_path + distributed_runtime = _distributed_runtime_config(case) + cmd = [python, "-c", script] + if distributed_runtime: + _ensure_distributed_runtime_env( + case, ctx, env, rendezvous_suffix=f".debug_{phase}") + extra_env = distributed_runtime.get("env", {}) + if isinstance(extra_env, dict): + env.update({str(k): str(v) for k, v in extra_env.items()}) + cmd = _wrap_distributed_command(cmd, case, env) + t0 = time.monotonic() + try: + result = subprocess.run( + cmd, + capture_output=True, text=True, timeout=600, env=env, + ) + except subprocess.TimeoutExpired: + elapsed = time.monotonic() - t0 + return None, elapsed, {"error": "timeout", "phase": phase} + except Exception as e: + elapsed = time.monotonic() - t0 + return None, elapsed, {"error": str(e), "phase": phase} + elapsed = time.monotonic() - t0 + + meta: dict = { + "returncode": result.returncode, + "command": cmd, + "stdout": result.stdout, + "stderr": result.stderr, + "phase": phase, + } + parse_stdout = ( + _extract_rank_zero_stdout(result.stdout) + if distributed_runtime + else result.stdout + ) + if distributed_runtime: + meta["distributed_runtime"] = distributed_runtime + meta["rank_zero_stdout"] = parse_stdout + meta["stderr_without_mpi_tags"] = _strip_mpi_stream_tags(result.stderr) + for line in parse_stdout.splitlines(): + if line.startswith("TRTMC_DEBUG_META "): + try: + parsed = json.loads(line[len("TRTMC_DEBUG_META "):]) + if isinstance(parsed, dict): + meta.update(parsed) + except json.JSONDecodeError: + meta["debug_meta_parse_error"] = line + if result.returncode != 0: + truncated, log_path = save_full_stderr( + result.stderr, ctx.artifacts_dir or "", + f"debug_runner_{phase}", case.name) + meta["stderr_truncated"] = truncated + if log_path: + meta["stderr_log"] = log_path + logger.warning( + "Debug runner (%s) failed for %s (rc=%d): %s", + phase, case.name, result.returncode, result.stderr[-500:] + ) + return None, elapsed, meta + + if not Path(logits_path).is_file(): + meta["error"] = "logits file not created" + return None, elapsed, meta + + return logits_path, elapsed, meta + + +plugin = TextGenerationCausalRunner() diff --git a/tests/e2e/models/smollm3/manifests/smollm3-3b.json b/tests/e2e/models/smollm3/manifests/smollm3-3b.json new file mode 100644 index 0000000000..4e0660a13e --- /dev/null +++ b/tests/e2e/models/smollm3/manifests/smollm3-3b.json @@ -0,0 +1,24 @@ +{ + "name": "smollm3-3b", + "hf_id": "HuggingFaceTB/SmolLM3-3B", + "hf_revision": "a07cc9a04f16550a088caea529712d1d335b0ac1", + "bundle": "smollm3-3b.bundle", + "family": "smollm3", + "runtime_strategy": "smollm3_decoder_kv_cache", + "task_strategy": "text_generation_causal", + "e2e_parallel_resource": "exclusive_gpu", + "precision": "bf16", + "max_cache_length": 256, + "trust_remote_code": false, + "testcases": [ + { + "name": "smollm3-3b", + "trace_id": "IT-E2E-SMOLLM3-01", + "reference_family": "causal_base_continuation", + "user_contract": "continuation_parity", + "reference_precision": "fp32", + "prompt": "The quick brown fox jumps over the lazy dog. Once upon a time", + "max_new_tokens": 20 + } + ] +} diff --git a/tests/e2e/models/smollm3/runner.py b/tests/e2e/models/smollm3/runner.py new file mode 100644 index 0000000000..34caf7b129 --- /dev/null +++ b/tests/e2e/models/smollm3/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 smollm3 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/smollm3/test_smollm3_build_contract.py b/tests/e2e/models/smollm3/test_smollm3_build_contract.py new file mode 100644 index 0000000000..b937df1b54 --- /dev/null +++ b/tests/e2e/models/smollm3/test_smollm3_build_contract.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Family-owned CI build contracts for SmolLM3 models.""" + +from __future__ import annotations + +import json +from pathlib import Path + +from tests.e2e_harness.manifest_loader import load_manifest + +MANIFEST = Path(__file__).parent / "manifests" / "smollm3-3b.json" + + +def _manifest() -> dict: + return json.loads(MANIFEST.read_text(encoding="utf-8")) + + +def test_manifest_pins_an_immutable_revision() -> None: + assert load_manifest(MANIFEST).hf_revision == ( + "a07cc9a04f16550a088caea529712d1d335b0ac1" + ) + + +def test_manifest_declares_the_family_runtime_contract() -> None: + manifest = _manifest() + assert manifest["family"] == "smollm3" + assert manifest["runtime_strategy"] == "smollm3_decoder_kv_cache" + assert manifest["task_strategy"] == "text_generation_causal" + + +def test_manifest_builds_bf16_for_the_native_kv_path() -> None: + # build_routing accepts only BF16 for the native KV decoder, and the + # plugin's default_build_precision returns bf16 once the architecture + # qualifies, so the declared precision has to agree with both. + manifest = _manifest() + assert manifest["precision"] == "bf16" + assert load_manifest(MANIFEST).metadata["precision"] == "bf16" + + +def test_manifest_reserves_an_exclusive_gpu() -> None: + assert _manifest()["e2e_parallel_resource"] == "exclusive_gpu" + + +def test_manifest_uses_the_hf_transformers_oracle() -> None: + case = load_manifest(MANIFEST) + assert case.reference_backend == "hf_transformers" + assert case.oracle_level == "L1_external_reference" + assert case.reference_family == "causal_base_continuation" + assert case.user_contract == "continuation_parity" + + +def test_manifest_needs_no_remote_code() -> None: + # SmolLM3 is native in transformers; the bundle must not depend on + # trust_remote_code. + assert _manifest()["trust_remote_code"] is False diff --git a/tests/e2e/models/smollm3/test_smollm3_builder_engine.py b/tests/e2e/models/smollm3/test_smollm3_builder_engine.py new file mode 100644 index 0000000000..ccdf9dfad6 --- /dev/null +++ b/tests/e2e/models/smollm3/test_smollm3_builder_engine.py @@ -0,0 +1,195 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Engine tests for the SmolLM3 family plugin. + +Trace: ARCH-FAM-001, UD-FAM-SMOLLM3-01 +Intent: Validate the SmolLM3 family plugin weight loading and standard decoder key mapping with RMSNorm and SwiGLU MLP. +Preconditions: safetensors and tensorrt_model_connect are importable; TRT+GPU required for engine build tests. +Postconditions: All standard decoder weight keys are present with correct shapes and the engine builds successfully. +""" +from __future__ import annotations + +import pytest + +from tests.builder.family_plugin_tester import FamilyPluginTester, TinyModelSpec +from tests.builder.family_plugin_test_mixin import ( + FamilyPluginTestMixin, + requires_trt, +) + + +class SmolLM3PluginTester(FamilyPluginTester): + plugin_module = "tensorrt_model_connect.families.smollm3" + model_type = "smollm3" + + +class NativeSmolLM3PluginTester(SmolLM3PluginTester): + """Smallest production-shaped dense SmolLM3 accepted by native attention.""" + + spec = TinyModelSpec( + vocab_size=32, + hidden_size=128, + intermediate_size=256, + num_hidden_layers=1, + num_attention_heads=1, + num_key_value_heads=1, + head_dim=128, + max_position_embeddings=128, + max_cache_length=128, + ) + + def get_config_dict(self) -> dict: + config = super().get_config_dict() + config.update( + architectures=["SmolLM3ForCausalLM"], + hidden_act="silu", + ) + return config + + +def _deserialize(plan: bytes): + import tensorrt as trt + + return trt.Runtime(trt.Logger(trt.Logger.WARNING)).deserialize_cuda_engine(plan) + + +def _io_names(engine) -> tuple[set[str], set[str]]: + import tensorrt as trt + + inputs: set[str] = set() + outputs: set[str] = set() + for index in range(engine.num_io_tensors): + name = engine.get_tensor_name(index) + target = ( + inputs + if engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT + else outputs + ) + target.add(name) + return inputs, outputs + + +class TestSmolLM3Engine(FamilyPluginTestMixin): + tester_class = SmolLM3PluginTester + + @staticmethod + def _build_legacy_engine(tester, tmp_path) -> bytes: + from tensorrt_model_connect.families.smollm3.standard_decoder_builder import ( + build_standard_decoder_engine, + ) + + config, weights, _ = tester.prepare_config_and_weights(tmp_path) + return build_standard_decoder_engine( + config, + weights, + tester.spec.max_cache_length, + precision="fp32", + verbose=False, + ) + + @requires_trt + def test_build_engine_succeeds(self, tester, tmp_path): + """Keep the generic dense-mask SmolLM3 builder smoke covered.""" + plan = self._build_legacy_engine(tester, tmp_path) + assert isinstance(plan, bytes) + assert plan + + @requires_trt + def test_engine_io_tensor_names(self, tester, tmp_path): + """Keep the legacy Python-builder/C++ tensor naming contract covered.""" + engine = _deserialize(self._build_legacy_engine(tester, tmp_path)) + assert engine is not None + inputs, outputs = _io_names(engine) + assert inputs == tester.expected_engine_input_names() + assert outputs == tester.expected_engine_output_names() + + @requires_trt + def test_engine_logits_output_shape(self, tester, tmp_path): + """Keep the legacy single-row logits contract covered.""" + engine = _deserialize(self._build_legacy_engine(tester, tmp_path)) + assert engine is not None + assert tuple(engine.get_tensor_shape("logits")) == ( + 1, + tester.spec.vocab_size, + ) + + @pytest.mark.parametrize("role", ["prefill", "decode"]) + @requires_trt + def test_native_split_role_engine_contract(self, tmp_path, role): + """Qualified BF16 split engines expose full-capacity aliased KV state.""" + import tensorrt as trt + + tester = NativeSmolLM3PluginTester() + config, weights, _ = tester.prepare_config_and_weights(tmp_path) + config.raw["_decoder_engine_role"] = role + + plan = tester.get_plugin().build_engine( + config, + weights, + tester.spec.max_cache_length, + precision="bf16", + verbose=False, + ) + engine = _deserialize(plan) + assert engine is not None + assert engine.num_optimization_profiles == 1 + + inputs, outputs = _io_names(engine) + assert "attention_mask" not in inputs + assert {"cache_write_indices", "key_value_lengths"} <= inputs + assert {"cache_k_0", "cache_v_0"} <= inputs + assert {"present_k_0", "present_v_0"} <= outputs + + cache_shape = ( + 1, + tester.spec.num_key_value_heads, + tester.spec.max_cache_length, + tester.spec.head_dim, + ) + for stem in ("k", "v"): + cache = f"cache_{stem}_0" + present = f"present_{stem}_0" + assert tuple(engine.get_tensor_shape(cache)) == cache_shape + assert tuple(engine.get_tensor_shape(present)) == cache_shape + assert engine.get_tensor_dtype(cache) == trt.bfloat16 + assert engine.get_aliased_input_tensor(present) == cache + + profile = engine.get_tensor_profile_shape("token_id", 0) + assert tuple(profile[0]) == (1,) + assert tuple(profile[2]) == ( + min(tester.spec.max_cache_length, 64) if role == "prefill" else 1, + ) + + @requires_trt + def test_native_decode_records_explicit_attention_graph(self, tmp_path): + """Native decode uses primitives and does not expose the old KVL Recipe.""" + from tensorrt_model_connect.tvm_ffi import graph_build + from tensorrt_model_connect.tvm_ffi.graph_patch import load_snapshot + + tester = NativeSmolLM3PluginTester() + config, weights, _ = tester.prepare_config_and_weights(tmp_path) + config.raw["_decoder_engine_role"] = "decode" + snapshot_path = tmp_path / "smollm3-decode.graph.json" + + with pytest.raises(graph_build.GraphInspectionComplete): + with graph_build.inspect_graph( + snapshot_path, + engine_role="decode", + metadata={}, + ): + with graph_build.engine_role("decode"): + tester.get_plugin().build_engine( + config, + weights, + tester.spec.max_cache_length, + precision="bf16", + verbose=False, + ) + + snapshot = load_snapshot(snapshot_path) + assert snapshot.metadata.get("graph_recipes", []) == [] + operations = [node.op for node in snapshot.nodes] + assert sum("MATRIX_MULTIPLY" in operation for operation in operations) >= 2 + assert any("SOFTMAX" in operation for operation in operations) + assert not any(operation.endswith("ATTENTION") for operation in operations) diff --git a/tests/e2e/models/smollm3/test_smollm3_e2e.py b/tests/e2e/models/smollm3/test_smollm3_e2e.py new file mode 100644 index 0000000000..cfc0856a03 --- /dev/null +++ b/tests/e2e/models/smollm3/test_smollm3_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 smollm3 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/smollm3/test_smollm3_family_plugin_weights.py b/tests/e2e/models/smollm3/test_smollm3_family_plugin_weights.py new file mode 100644 index 0000000000..f91bd8203d --- /dev/null +++ b/tests/e2e/models/smollm3/test_smollm3_family_plugin_weights.py @@ -0,0 +1,125 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Family-owned plugin weight tests. + +Concrete load_weights behavior belongs beside the model family it validates. +Shared test code is limited to filesystem and serialization helpers. +""" + +from __future__ import annotations + + +import numpy as np + +from tests.builder.family_plugin_test_support import ( + ModelConfig, + _rand, + _write_config, + _write_safetensors, +) + + +def _make_standard_decoder_tensors(vocab, hidden, layers, heads, kv_heads, mlp): + head_dim = hidden // heads + kv_hidden = kv_heads * head_dim + tensors = {} + tensors["model.embed_tokens.weight"] = _rand(vocab, hidden) + for i in range(layers): + prefix = f"model.layers.{i}" + tensors[f"{prefix}.input_layernorm.weight"] = _rand(hidden) + tensors[f"{prefix}.post_attention_layernorm.weight"] = _rand(hidden) + tensors[f"{prefix}.self_attn.q_proj.weight"] = _rand(hidden, hidden) + tensors[f"{prefix}.self_attn.k_proj.weight"] = _rand(kv_hidden, hidden) + tensors[f"{prefix}.self_attn.v_proj.weight"] = _rand(kv_hidden, hidden) + tensors[f"{prefix}.self_attn.o_proj.weight"] = _rand(hidden, hidden) + tensors[f"{prefix}.mlp.gate_proj.weight"] = _rand(mlp, hidden) + tensors[f"{prefix}.mlp.up_proj.weight"] = _rand(mlp, hidden) + tensors[f"{prefix}.mlp.down_proj.weight"] = _rand(hidden, mlp) + tensors["model.norm.weight"] = _rand(hidden) + tensors["lm_head.weight"] = _rand(vocab, hidden) + return tensors + + +class TestSmolLM3Plugin: + VOCAB, HIDDEN, LAYERS, HEADS, KV_HEADS, MLP = 32, 16, 2, 4, 2, 32 + + def test_selected_fp32_layers_use_single_engine_layout(self): + from tensorrt_model_connect.families.smollm3 import plugin + + config = ModelConfig( + hidden_size=self.HIDDEN, + vocab_size=self.VOCAB, + num_hidden_layers=self.LAYERS, + num_attention_heads=self.HEADS, + num_key_value_heads=self.KV_HEADS, + ) + assert plugin.supports_split_decoder_roles(config) + + config.raw["_fp32_layers"] = [1] + assert not plugin.supports_split_decoder_roles(config) + + def test_load_weights(self, tmp_path): + """SmolLM3 uses load_standard_weights — verify compact GQA K/V.""" + from tensorrt_model_connect.families.smollm3 import plugin + + head_dim = self.HIDDEN // self.HEADS # 4 + kv_hidden = self.KV_HEADS * head_dim # 8 + config = { + "model_type": "smollm3", + "vocab_size": self.VOCAB, + "hidden_size": self.HIDDEN, + "num_hidden_layers": self.LAYERS, + "num_attention_heads": self.HEADS, + "num_key_value_heads": self.KV_HEADS, + } + tensors = _make_standard_decoder_tensors( + self.VOCAB, self.HIDDEN, self.LAYERS, self.HEADS, self.KV_HEADS, + self.MLP) + _write_config(tmp_path, config) + _write_safetensors(tmp_path, tensors) + + cfg = ModelConfig.from_dir(tmp_path) + weights = plugin.load_weights(str(tmp_path), cfg) + + # K/V stay compact at [hidden, kv_hidden]. + for i in range(self.LAYERS): + assert weights[f"layer.{i}.w_k"].shape == ( + self.HIDDEN, kv_hidden) + assert weights[f"layer.{i}.w_v"].shape == ( + self.HIDDEN, kv_hidden) + + cfg.raw["_fp32_layers"] = [1] + mixed_weights = plugin.load_weights( + str(tmp_path), cfg, precision="fp16") + assert mixed_weights["embedding"].dtype == np.float16 + assert mixed_weights["layer.0.w_q"].dtype == np.float16 + assert mixed_weights["layer.1.w_q"].dtype == np.float32 + + def test_tied_embeddings(self, tmp_path): + """When lm_head.weight is missing, w_out = transposed embedding.""" + from tensorrt_model_connect.families.smollm3 import plugin + + config = { + "model_type": "smollm3", + "vocab_size": self.VOCAB, + "hidden_size": self.HIDDEN, + "num_hidden_layers": 1, + "num_attention_heads": self.HEADS, + "num_key_value_heads": self.KV_HEADS, + "tie_word_embeddings": True, + } + tensors = _make_standard_decoder_tensors( + self.VOCAB, self.HIDDEN, 1, self.HEADS, self.KV_HEADS, self.MLP) + # Remove lm_head to test tied embedding fallback + del tensors["lm_head.weight"] + _write_config(tmp_path, config) + _write_safetensors(tmp_path, tensors) + + cfg = ModelConfig.from_dir(tmp_path) + weights = plugin.load_weights(str(tmp_path), cfg) + + assert weights["w_out"].shape == (self.HIDDEN, self.VOCAB) + embedding = tensors["model.embed_tokens.weight"] + np.testing.assert_allclose( + weights["w_out"], embedding.T, atol=1e-6) diff --git a/tests/e2e/models/smollm3/test_smollm3_native_kv_routing.py b/tests/e2e/models/smollm3/test_smollm3_native_kv_routing.py new file mode 100644 index 0000000000..ec598d9d67 --- /dev/null +++ b/tests/e2e/models/smollm3/test_smollm3_native_kv_routing.py @@ -0,0 +1,842 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU-only contract tests for SmolLM3's TensorRT native KV path.""" + +from __future__ import annotations + +from dataclasses import dataclass +import importlib + +import pytest + +from tensorrt_model_connect.families.smollm3.build_routing import ( + native_kv_architecture_capability, + native_kv_build_capability, + native_kv_cache_geometry, + prefer_native_default, + resolved_head_dim, +) +from tensorrt_model_connect.families.smollm3.config import ModelConfig +from tensorrt_model_connect.families.smollm3.native_kv_contract import ( + validate_native_kv_weights, +) + + +_LLAMA3_ROPE = { + "rope_type": "llama3", + "factor": 8.0, + "low_freq_factor": 1.0, + "high_freq_factor": 4.0, + "original_max_position_embeddings": 8192, +} + + +def _config( + *, + raw_updates: dict | None = None, + llama3_rope: bool = True, + **overrides, +) -> ModelConfig: + values = { + "model_type": "smollm3", + "architectures": ["SmolLM3ForCausalLM"], + "vocab_size": 128256, + "hidden_size": 4096, + "intermediate_size": 14336, + "num_hidden_layers": 32, + "num_attention_heads": 32, + "num_key_value_heads": 8, + "rms_norm_eps": 1e-5, + "rope_theta": 500_000.0, + "max_position_embeddings": 131072, + "hidden_act": "silu", + "_head_dim": 0, + } + values.update(overrides) + raw = { + "_decoder_engine_layout": "split", + "rope_scaling": dict(_LLAMA3_ROPE) if llama3_rope else None, + } + raw.update(raw_updates or {}) + values["raw"] = raw + return ModelConfig(**values) + + +@pytest.mark.parametrize( + ( + "hidden", + "mlp", + "layers", + "heads", + "kv_heads", + "context", + ), + [ + (4096, 14336, 32, 32, 8, 131072), + (5120, 13824, 40, 40, 40, 4096), + (8192, 28672, 80, 64, 8, 131072), + ], + ids=("smollm3-8b-shape", "smollm3-13b-shape", "smollm3-70b-shape"), +) +def test_dense_smollm3_sizes_share_one_native_contract( + hidden, mlp, layers, heads, kv_heads, context, +): + config = _config( + hidden_size=hidden, + intermediate_size=mlp, + num_hidden_layers=layers, + num_attention_heads=heads, + num_key_value_heads=kv_heads, + max_position_embeddings=context, + ) + + architecture = native_kv_architecture_capability(config) + build = native_kv_build_capability(config) + row_bytes, cache_bytes = native_kv_cache_geometry(config, context) + + assert architecture.eligible, architecture.reason + assert build.eligible, build.reason + assert prefer_native_default(config) + assert resolved_head_dim(config) == 128 + assert row_bytes == 2 * layers * kv_heads * 128 * 2 + assert cache_bytes == context * row_bytes + + +def test_route_uses_architecture_not_checkpoint_identity(): + config = _config( + raw_updates={ + "_model_dir": "/models/renamed-checkpoint", + "name_or_path": "any-owner/any-smollm3", + "checkpoint_sha256": "a" * 64, + } + ) + + assert native_kv_architecture_capability(config).eligible + assert prefer_native_default(config) + + +def test_explicit_head_dim_is_supported_when_hidden_width_is_decoupled(): + config = _config( + hidden_size=3072, + num_attention_heads=32, + _head_dim=128, + ) + + assert resolved_head_dim(config) == 128 + assert native_kv_architecture_capability(config).eligible + + +@pytest.mark.parametrize( + ("overrides", "raw_updates", "reason"), + [ + ({"model_type": "smollm34"}, {}, "model_type"), + ({"architectures": ["OtherForCausalLM"]}, {}, "architectures"), + ({"hidden_size": 4100}, {}, "divisible"), + ({"_head_dim": 64}, {}, "head_dim=128"), + ({"num_key_value_heads": 6}, {}, "divisible"), + ({"hidden_act": "gelu"}, {}, "hidden_act"), + ({}, {"sliding_window": 4096}, "unsupported SmolLM3 fields"), + ({}, {"num_experts": 8}, "unsupported SmolLM3 fields"), + ( + {}, + {"layer_types": ["full_attention", "linear_attention"]}, + "hybrid", + ), + ( + {}, + {"rope_scaling": {"rope_type": "linear", "factor": 2.0}}, + "rope_type", + ), + ], +) +def test_architecture_variants_fail_closed(overrides, raw_updates, reason): + decision = native_kv_architecture_capability( + _config(raw_updates=raw_updates, **overrides) + ) + + assert decision.applicable + assert not decision.eligible + assert reason in decision.reason + + +@pytest.mark.parametrize( + ("kwargs", "raw_updates", "reason"), + [ + ({"precision": "fp16"}, {}, "BF16"), + ({"max_cache_length": 131071}, {}, "max_cache_length"), + ({"parallel_enabled": True}, {}, "tensor parallel"), + ({"dynamic_kv_cache": True}, {}, "fixed physical"), + ({"quantized": True}, {}, "quantized"), + ({"debug_layer_outputs": True}, {}, "debug"), + ({}, {"_fp32_layers": ["layer.0"]}, "FP32 layer"), + ({}, {"_decoder_engine_layout": "dual_profile"}, "split"), + ({}, {"_rtx_build_requested": True}, "standard TensorRT"), + ], +) +def test_unqualified_build_modes_fail_closed(kwargs, raw_updates, reason): + decision = native_kv_build_capability( + _config(raw_updates=raw_updates), + **kwargs, + ) + + assert not decision.eligible + assert reason in decision.reason + + +@dataclass +class _Tensor: + shape: tuple[int, ...] + + +def _small_config(*, role: str = "prefill") -> ModelConfig: + return _config( + vocab_size=32, + hidden_size=128, + intermediate_size=256, + num_hidden_layers=1, + num_attention_heads=1, + num_key_value_heads=1, + max_position_embeddings=256, + llama3_rope=False, + raw_updates={"_decoder_engine_role": role}, + ) + + +def _weights(config: ModelConfig) -> dict[str, object]: + hidden = config.hidden_size + attention = config.num_attention_heads * 128 + kv_attention = config.num_key_value_heads * 128 + mlp = config.intermediate_size + weights: dict[str, object] = { + "embedding": _Tensor((config.vocab_size, hidden)), + "final_norm": _Tensor((hidden,)), + "w_out": _Tensor((hidden, config.vocab_size)), + "_attention_size": attention, + "_kv_attention_size": kv_attention, + "_mlp_size": mlp, + } + for name, shape in ( + ("input_norm", (hidden,)), + ("w_q", (hidden, attention)), + ("w_k", (hidden, kv_attention)), + ("w_v", (hidden, kv_attention)), + ("w_o", (attention, hidden)), + ("post_attn_norm", (hidden,)), + ("w_gate", (hidden, mlp)), + ("w_up", (hidden, mlp)), + ("w_down", (mlp, hidden)), + ): + weights[f"layer.0.{name}"] = _Tensor(shape) + return weights + + +def test_weight_contract_rejects_missing_shape_and_bias(): + config = _small_config() + weights = _weights(config) + validate_native_kv_weights(config, weights) + + missing = dict(weights) + missing.pop("layer.0.w_k") + with pytest.raises(ValueError, match="missing.*w_k"): + validate_native_kv_weights(config, missing) + + wrong_shape = dict(weights) + wrong_shape["layer.0.w_q"] = _Tensor((127, 128)) + with pytest.raises(ValueError, match="must have shape"): + validate_native_kv_weights(config, wrong_shape) + + biased = dict(weights) + biased["layer.0.q_bias"] = _Tensor((128,)) + with pytest.raises(ValueError, match="bias"): + validate_native_kv_weights(config, biased) + + +def test_plugin_builds_the_requested_split_role_directly(monkeypatch): + pytest.importorskip("tensorrt") + plugin_module = importlib.import_module( + "tensorrt_model_connect.families.smollm3.plugin" + ) + + config = _small_config(role="prefill") + captured: dict[str, object] = {} + + def _build(*args, **kwargs): + captured.update(args=args, kwargs=kwargs) + return b"plan" + + monkeypatch.setattr( + plugin_module, + "build_dual_profile_decoder_engine", + _build, + ) + + result = plugin_module.plugin.build_engine( + config, + _weights(config), + 256, + precision="bf16", + ) + + assert result == b"plan" + assert captured["kwargs"]["profile_mode"] == "prefill" + assert captured["kwargs"]["native_kv_cache"] is True + assert plugin_module.plugin.get_bundle_config_overrides(config) == { + "native_kv_contract_version": 1, + "native_kv_cache": True, + } + + +def test_plugin_falls_back_for_explicit_legacy_build_options(monkeypatch): + pytest.importorskip("tensorrt") + plugin_module = importlib.import_module( + "tensorrt_model_connect.families.smollm3.plugin" + ) + + config = _small_config(role="decode") + config.raw["_native_kv_cache_metadata"] = {"stale": True} + quant_ctx = object() + captured: dict[str, object] = {} + + def _build(*args, **kwargs): + captured.update(args=args, kwargs=kwargs) + return b"legacy-plan" + + monkeypatch.setattr( + plugin_module, + "build_standard_decoder_engine", + _build, + ) + + result = plugin_module.plugin.build_engine( + config, + _weights(config), + 128, + precision="fp16", + quant_ctx=quant_ctx, + ) + + assert result == b"legacy-plan" + assert captured["args"][2] == 128 + assert captured["kwargs"]["precision"] == "fp16" + assert captured["kwargs"]["quant_ctx"] is quant_ctx + assert plugin_module.plugin.get_bundle_config_overrides(config) is None + + +def test_dynamic_kv_dual_profile_dispatches_bucket_rows(monkeypatch): + pytest.importorskip("tensorrt") + builder_module = importlib.import_module( + "tensorrt_model_connect.families.smollm3.standard_decoder_builder" + ) + config = _small_config(role="dual_profile") + config.raw["dynamic_kv_cache"] = True + config.raw["_dynamic_kv_profile_rows"] = [256, 131072] + captured: dict[str, object] = {} + + def _build(*args, **kwargs): + captured.update(args=args, kwargs=kwargs) + return b"dynamic-dual-profile-plan" + + monkeypatch.setattr( + builder_module, + "build_dual_profile_decoder_engine", + _build, + ) + + plan = builder_module.build_standard_decoder_engine( + config, + _weights(config), + 131072, + precision="fp16", + ) + + assert plan == b"dynamic-dual-profile-plan" + assert captured["args"][2] == 131072 + assert captured["kwargs"]["precision"] == "fp16" + assert captured["kwargs"]["dynamic_kv_profile_rows"] == [256, 131072] + assert captured["kwargs"]["profile_mode"] == "dual_profile" + + config.raw.pop("_dynamic_kv_profile_rows") + captured.clear() + plan = builder_module.build_standard_decoder_engine( + config, + _weights(config), + 131072, + precision="fp16", + ) + + assert plan == b"dynamic-dual-profile-plan" + assert captured["kwargs"]["dynamic_kv_profile_rows"] == [131072] + + +def test_plugin_falls_back_outside_the_native_architecture_contract( + monkeypatch, +): + pytest.importorskip("tensorrt") + plugin_module = importlib.import_module( + "tensorrt_model_connect.families.smollm3.plugin" + ) + config = _small_config() + config._head_dim = 64 + captured: dict[str, object] = {} + + def _build(*args, **kwargs): + captured.update(args=args, kwargs=kwargs) + return b"legacy-plan" + + monkeypatch.setattr( + plugin_module, + "build_standard_decoder_engine", + _build, + ) + + assert not prefer_native_default(config) + assert plugin_module.plugin.default_build_precision(config) == "fp32" + assert plugin_module.plugin.default_max_cache_length(config) == 256 + assert plugin_module.plugin.build_engine( + config, + _weights(config), + 128, + precision="fp16", + ) == b"legacy-plan" + assert captured["args"][2] == 128 + assert captured["kwargs"]["precision"] == "fp16" + + +def _shared_build_config(**raw_updates): + """Build the config the engine builder actually constructs. + + ``engine_builder`` resolves a checkpoint through the shared + ``tensorrt_model_connect.config.ModelConfig``, never this family's + dataclass, so anything the builders read has to work on that object. The + typed fields mirror what ``from_dir`` fills in for SmolLM3-3B. + """ + from tensorrt_model_connect.config import ModelConfig as SharedModelConfig + + raw = { + "_decoder_engine_layout": "split", + "no_rope_layer_interval": 4, + "rope_scaling": None, + } + raw.update(raw_updates) + return SharedModelConfig( + model_type="smollm3", + architectures=["SmolLM3ForCausalLM"], + vocab_size=128256, + hidden_size=2048, + intermediate_size=11008, + num_hidden_layers=36, + num_attention_heads=16, + num_key_value_heads=4, + rms_norm_eps=1e-5, + rope_theta=5_000_000.0, + max_position_embeddings=65536, + hidden_act="silu", + tie_word_embeddings=True, + raw=raw, + ) + + +def test_rope_layer_schedule_resolves_on_the_shared_build_config(): + from tensorrt_model_connect.families.smollm3.config import ( + resolve_rope_layer_schedule, + ) + + schedule = resolve_rope_layer_schedule(_shared_build_config()) + + assert len(schedule) == 36 + assert [index for index, uses in enumerate(schedule) if not uses] == [ + 3, 7, 11, 15, 19, 23, 27, 31, 35 + ] + assert schedule == ModelConfig( + model_type="smollm3", + num_hidden_layers=36, + raw=dict(_shared_build_config().raw), + ).rope_layer_schedule(), "family and shared configs must agree" + + +def test_published_no_rope_layers_wins_over_the_interval(): + from tensorrt_model_connect.families.smollm3.config import ( + resolve_rope_layer_schedule, + ) + + published = [1] * 36 + published[5] = 0 + schedule = resolve_rope_layer_schedule( + _shared_build_config(no_rope_layers=published) + ) + + assert [index for index, uses in enumerate(schedule) if not uses] == [5] + + +@pytest.mark.parametrize( + "raw_updates, fragment", + [ + ({"no_rope_layer_interval": 0}, "no_rope_layer_interval must be positive"), + ({"no_rope_layers": [1, 1, 1]}, "no_rope_layers must be a sequence"), + ], +) +def test_malformed_schedule_is_rejected_on_the_shared_build_config( + raw_updates, fragment +): + from tensorrt_model_connect.families.smollm3.config import ( + resolve_rope_layer_schedule, + ) + + with pytest.raises(ValueError, match=fragment): + resolve_rope_layer_schedule(_shared_build_config(**raw_updates)) + + +def test_routing_rejects_a_malformed_schedule_on_the_shared_build_config(): + """Routing must judge the schedule on the config the build path carries. + + Resolving through a family-local method left this check silently inert for + the shared config, so a malformed schedule routed as eligible and only + surfaced once the graph builder ran. + """ + capability = native_kv_architecture_capability( + _shared_build_config(no_rope_layer_interval=0) + ) + + assert not capability.eligible + assert any("no_rope_layer_interval must be positive" in reason + for reason in capability.reason.split("; ")) + + +def test_routing_still_accepts_a_well_formed_schedule(): + assert native_kv_architecture_capability(_shared_build_config()).eligible + + +class _FakeTensor: + """Stand-in for an ITensor; the graph is never realized in this test.""" + + def __init__(self, name="t", shape=(1, 1, 64)): + self.name = name + self.shape = shape + self.dtype = None + + def __getattr__(self, _name): + return None + + +class _FakeLayer: + def get_output(self, _index): + return _FakeTensor("out") + + def __getattr__(self, _name): + return lambda *args, **kwargs: None + + def __setattr__(self, _name, _value): + pass + + +class _FakeNetwork: + """Accepts any add_* call and returns a layer with one output.""" + + def __getattr__(self, name): + if name.startswith("add_"): + return lambda *args, **kwargs: _FakeLayer() + raise AttributeError(name) + + +def _nope_block_weights(prefix, hidden, attention): + import numpy as np + + return { + f"{prefix}.input_norm": np.ones(hidden, dtype=np.float32), + f"{prefix}.w_q": np.zeros((hidden, attention), dtype=np.float32), + f"{prefix}.w_k": np.zeros((hidden, attention), dtype=np.float32), + f"{prefix}.w_v": np.zeros((hidden, attention), dtype=np.float32), + f"{prefix}.w_o": np.zeros((attention, hidden), dtype=np.float32), + } + + +def _count_rope_insertions(monkeypatch, *, apply_rope): + """Drive the attention block and count RoPE layer insertions. + + Spies on ``add_apply_rope_native``, which is what actually puts an + IRotaryEmbeddingLayer into the graph, so this observes the gate rather than + the source that contains it. + """ + from tensorrt_model_connect.families.smollm3 import graph_blocks, graph_ops + + hidden = attention = 64 + calls: list[tuple] = [] + monkeypatch.setattr( + graph_ops, + "add_apply_rope_native", + lambda *args, **kwargs: calls.append(args) or _FakeTensor("roped"), + ) + graph_blocks.add_attention_block( + _FakeNetwork(), + _FakeTensor("hidden"), + _FakeTensor("cache_k"), + _FakeTensor("cache_v"), + _FakeTensor("mask"), + _FakeTensor("position"), + weights=_nope_block_weights("layer.0", hidden, attention), + prefix="layer.0", + hidden_size=hidden, + attention_size=attention, + num_heads=2, + head_dim=32, + max_cache_length=16, + eps_tensor=_FakeTensor("eps"), + num_kv_heads=2, + kv_attention_size=attention, + apply_rope=apply_rope, + cos_half_tensor=_FakeTensor("cos"), + sin_half_tensor=_FakeTensor("sin"), + ) + return calls + + +def test_rope_layer_inserts_rotary_embedding_for_query_and_key(monkeypatch): + assert len(_count_rope_insertions(monkeypatch, apply_rope=True)) == 2 + + +def test_nope_layer_inserts_no_rotary_embedding_at_all(monkeypatch): + assert _count_rope_insertions(monkeypatch, apply_rope=False) == [] + + +def _rope_guard_sources(): + """Return, per builder, the schedule subscripts guarding RoPE insertion. + + The builders need TensorRT to run, so this reads their parsed source rather + than driving them. It checks the wiring the block-level tests above cannot + reach: that each builder selects the flag with its own layer loop variable + instead of a constant or the wrong index. + """ + import ast + import pathlib + + family = pathlib.Path( + "python/tensorrt_model_connect/families/smollm3" + ) + if not family.is_dir(): + import tensorrt_model_connect.families.smollm3 as package + + family = pathlib.Path(package.__file__).parent + + found = {} + for name in ("standard_decoder_builder.py", "dual_profile_decoder_builder.py"): + tree = ast.parse((family / name).read_text(encoding="utf-8")) + subscripts = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Subscript) + and isinstance(node.value, ast.Name) + and node.value.id == "rope_schedule" + ] + found[name] = [ + node.slice.id + for node in subscripts + if isinstance(node.slice, ast.Name) + ] + return found + + +def test_both_builders_select_the_flag_with_their_layer_index(): + guards = _rope_guard_sources() + + for name, indices in guards.items(): + assert indices, f"{name} does not subscript rope_schedule at all" + assert set(indices) == {"layer_idx"}, ( + f"{name} indexes rope_schedule with {sorted(set(indices))}, " + "which is not the layer loop variable" + ) + + +def test_standard_builder_forwards_the_flag_to_the_attention_block(): + """The standard builder reaches RoPE through ``apply_rope=``. + + A dropped keyword would silently restore RoPE on every NoPE layer, which + the block-level tests cannot see because they pass the flag directly. + """ + import ast + import pathlib + + import tensorrt_model_connect.families.smollm3 as package + + path = pathlib.Path(package.__file__).parent / "standard_decoder_builder.py" + tree = ast.parse(path.read_text(encoding="utf-8")) + forwarded = [ + keyword + for node in ast.walk(tree) + if isinstance(node, ast.Call) + for keyword in node.keywords + if keyword.arg == "apply_rope" + ] + + assert forwarded, "standard builder never passes apply_rope" + assert any( + isinstance(keyword.value, ast.Subscript) + and isinstance(keyword.value.value, ast.Name) + and keyword.value.value.id == "rope_schedule" + for keyword in forwarded + ), "apply_rope is passed but not from the resolved schedule" + + +def _published_checkpoint_config(): + """The published SmolLM3-3B configuration, field for field. + + Taken from config.json at the revision the manifest pins. It is spelled out + rather than trimmed so the routing contract is exercised against what users + actually download, including the fields this family does not read. + """ + from tensorrt_model_connect.config import ModelConfig as SharedModelConfig + + return SharedModelConfig( + model_type="smollm3", + architectures=["SmolLM3ForCausalLM"], + vocab_size=128256, + hidden_size=2048, + intermediate_size=11008, + num_hidden_layers=36, + num_attention_heads=16, + num_key_value_heads=4, + rms_norm_eps=1e-06, + rope_theta=5000000.0, + max_position_embeddings=65536, + hidden_act="silu", + tie_word_embeddings=True, + bos_token_id=128000, + eos_token_id=128012, + pad_token_id=128004, + raw={ + "attention_bias": False, + "attention_dropout": 0.0, + "layer_types": ["full_attention"] * 36, + "max_window_layers": 28, + "mlp_bias": False, + "no_rope_layer_interval": 4, + "no_rope_layers": [ + int((index + 1) % 4 != 0) for index in range(36) + ], + "pretraining_tp": 2, + "rope_scaling": None, + "sliding_window": None, + "torch_dtype": "bfloat16", + "use_cache": False, + "use_sliding_window": False, + }, + ) + + +def test_published_checkpoint_reaches_the_native_kv_path(): + """The checkpoint this family targets must route to its own runtime. + + The published config carries pretraining_tp=2, a field SmolLM3ForCausalLM + neither defines nor reads. Gating on it sent the default build of the only + supported checkpoint to the fallback decoder. + """ + config = _published_checkpoint_config() + + decision = native_kv_architecture_capability(config) + + assert decision.applicable + assert decision.eligible, decision.reason + assert prefer_native_default(config) + + +def test_default_build_of_the_published_checkpoint_is_native(): + """`trtmc build HuggingFaceTB/SmolLM3-3B` with no flags takes the native path.""" + import importlib + + plugin_module = importlib.import_module( + "tensorrt_model_connect.families.smollm3.plugin" + ) + config = _published_checkpoint_config() + + precision = plugin_module.plugin.default_build_precision(config) + cache_length = plugin_module.plugin.default_max_cache_length(config) + + assert precision == "bf16" + assert cache_length == config.max_position_embeddings == 65536 + + decision = native_kv_build_capability( + config, precision=precision, max_cache_length=cache_length + ) + assert decision.eligible, decision.reason + + +def _routing_loaded_as_production_does(): + """Load build_routing.py the way the family loader does. + + MODEL.toml points ``default_build_route`` at ``build_routing.py| + prefer_native_default``, and families/__init__.py loads that with + spec_from_file_location under a synthetic top-level name. The module + therefore has no package context at runtime, while the tests above import + it as part of the package, where a relative import would work. Load it + both ways so the difference cannot hide a failure again. + """ + import importlib.util + import pathlib + + import tensorrt_model_connect.families.smollm3 as package + + path = pathlib.Path(package.__file__).parent / "build_routing.py" + spec = importlib.util.spec_from_file_location( + "_trtmc_family_smollm3_build_routing", path + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_routing_works_without_package_context(): + routing = _routing_loaded_as_production_does() + config = _published_checkpoint_config() + + # prefer_native_default is the entry point MODEL.toml names. + assert routing.prefer_native_default(config) is True + assert routing.native_kv_architecture_capability(config).eligible + + +def test_no_runtime_name_is_imported_only_for_type_checking(): + """A name called at runtime must not come from a TYPE_CHECKING import. + + That block does not execute, so such a name raises NameError the first + time the function runs. The builders need TensorRT to run, so this is + checked on the parsed source rather than by driving them. + """ + import ast + import pathlib + + import tensorrt_model_connect.families.smollm3 as package + + family = pathlib.Path(package.__file__).parent + offenders = {} + for path in sorted(family.glob("*.py")): + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + + type_checking_only = set() + for node in ast.walk(tree): + if not isinstance(node, ast.If): + continue + try: + test_source = ast.unparse(node.test).strip() + except AttributeError: # pragma: no cover - Python < 3.9 + continue + if test_source not in ("TYPE_CHECKING", "typing.TYPE_CHECKING"): + continue + for sub in ast.walk(node): + if isinstance(sub, (ast.Import, ast.ImportFrom)): + for alias in sub.names: + type_checking_only.add( + alias.asname or alias.name.split(".")[0] + ) + + called = { + node.func.id + for node in ast.walk(tree) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) + } + overlap = type_checking_only & called + if overlap: + offenders[path.name] = sorted(overlap) + + assert not offenders, ( + "these names are called at runtime but imported only under " + f"TYPE_CHECKING: {offenders}" + ) diff --git a/tests/e2e/models/smollm3/test_smollm3_registry_contract.py b/tests/e2e/models/smollm3/test_smollm3_registry_contract.py new file mode 100644 index 0000000000..7d88a5f0e1 --- /dev/null +++ b/tests/e2e/models/smollm3/test_smollm3_registry_contract.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Family-owned registry contract tests.""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("tensorrt", reason="registry contract tests import plugin modules") + +from tensorrt_model_connect.families import find_plugin + + +def _plugin(model_type: str): + plugin = find_plugin(model_type) + assert plugin is not None + return plugin + +def test_runtime_strategy() -> None: + plugin = _plugin("smollm3") + assert getattr(plugin, "runtime_strategy", None) == "smollm3_decoder_kv_cache" + + +def test_no_embed_input() -> None: + plugin = _plugin("smollm3") + assert not getattr(plugin, "embed_input", False) diff --git a/tests/e2e/models/smollm3/test_smollm3_rope_scaling.py b/tests/e2e/models/smollm3/test_smollm3_rope_scaling.py new file mode 100644 index 0000000000..38ac73351a --- /dev/null +++ b/tests/e2e/models/smollm3/test_smollm3_rope_scaling.py @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""SmolLM3 YaRN RoPE scaling tests. + +SmolLM3-3B ships with ``rope_scaling: null`` and a 65536-token window. The +model card extends that to 128k by setting ``max_position_embeddings`` to +131072 and adding a YaRN block, which is the configuration these tests pin. + +The reference below is the exact float64 form of Hugging Face's +``_compute_yarn_parameters``. One detail is easy to miss and is asserted +separately: upstream folds ``attention_factor`` (``0.1 * ln(factor) + 1``) into +cos/sin inside the rotary embedding rather than applying it to attention +scores, so a table built without it is uniformly off by that factor. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from tensorrt_model_connect.families.smollm3 import graph_ops + + +YARN_SCALING = { + "rope_type": "yarn", + "factor": 2.0, + "original_max_position_embeddings": 65536, +} +HEAD_DIM = 128 +ROPE_THETA = 5000000.0 + + +def _yarn_reference(head_dim: int, rope_theta: float, factor: float, + original_context: float) -> tuple[np.ndarray, float]: + """Exact float64 YaRN inverse frequencies and attention factor.""" + extrapolation = rope_theta ** ( + -np.arange(0, head_dim, 2, dtype=np.float64) / head_dim + ) + half = head_dim // 2 + + def correction_dim(rotations: float) -> float: + return ( + head_dim + * np.log(original_context / (rotations * 2.0 * np.pi)) + / (2.0 * np.log(rope_theta)) + ) + + low = max(int(np.floor(correction_dim(32.0))), 0) + high = min(int(np.ceil(correction_dim(1.0))), half - 1) + ramp = np.clip( + (np.arange(half, dtype=np.float64) - low) / max(high - low, 1), 0.0, 1.0 + ) + inverse = extrapolation / factor * ramp + extrapolation * (1.0 - ramp) + return inverse, 0.1 * np.log(factor) + 1.0 + + +def test_half_dim_table_matches_yarn_reference_formula() -> None: + positions = 32 + inverse, attention_factor = _yarn_reference( + HEAD_DIM, ROPE_THETA, 2.0, 65536.0 + ) + angles = np.outer(np.arange(positions, dtype=np.float64), inverse) + + cosine = graph_ops.make_rope_table_half_dim( + positions, HEAD_DIM, ROPE_THETA, True, rope_scaling=YARN_SCALING + ) + sine = graph_ops.make_rope_table_half_dim( + positions, HEAD_DIM, ROPE_THETA, False, rope_scaling=YARN_SCALING + ) + + np.testing.assert_allclose( + cosine, np.cos(angles) * attention_factor, atol=1e-7 + ) + np.testing.assert_allclose( + sine, np.sin(angles) * attention_factor, atol=1e-7 + ) + + +def test_attention_factor_is_folded_into_the_table() -> None: + # Omitting it leaves every entry short by ~6.9% for factor=2.0, which is + # four orders of magnitude above float32 noise. + inverse, attention_factor = _yarn_reference( + HEAD_DIM, ROPE_THETA, 2.0, 65536.0 + ) + assert attention_factor == pytest.approx(1.0693147180559945) + + cosine = graph_ops.make_rope_table_half_dim( + 8, HEAD_DIM, ROPE_THETA, True, rope_scaling=YARN_SCALING + ) + unscaled = np.cos(np.outer(np.arange(8, dtype=np.float64), inverse)) + assert np.abs(cosine - unscaled).max() > 1e-3 + np.testing.assert_allclose(cosine, unscaled * attention_factor, atol=1e-7) + + +def test_explicit_attention_factor_overrides_the_derived_one() -> None: + scaling = dict(YARN_SCALING, attention_factor=1.0) + inverse, _ = _yarn_reference(HEAD_DIM, ROPE_THETA, 2.0, 65536.0) + cosine = graph_ops.make_rope_table_half_dim( + 8, HEAD_DIM, ROPE_THETA, True, rope_scaling=scaling + ) + expected = np.cos(np.outer(np.arange(8, dtype=np.float64), inverse)) + np.testing.assert_allclose(cosine, expected, atol=1e-7) + + +@pytest.mark.parametrize( + "override", + [ + {"factor": 0.0}, + {"original_max_position_embeddings": 0}, + {"beta_fast": 1.0, "beta_slow": 32.0}, + ], +) +def test_malformed_yarn_scaling_is_rejected(override) -> None: + with pytest.raises(ValueError): + graph_ops.make_rope_table_half_dim( + 8, HEAD_DIM, ROPE_THETA, True, + rope_scaling=dict(YARN_SCALING, **override), + ) + + +def test_unscaled_rope_is_unaffected_by_yarn_support() -> None: + positions = 16 + inverse = ROPE_THETA ** ( + -np.arange(0, HEAD_DIM, 2, dtype=np.float64) / HEAD_DIM + ) + angles = np.outer(np.arange(positions, dtype=np.float64), inverse) + cosine = graph_ops.make_rope_table_half_dim( + positions, HEAD_DIM, ROPE_THETA, True + ) + np.testing.assert_allclose(cosine, np.cos(angles), atol=1e-7) diff --git a/tests/e2e/models/smollm3/test_smollm3_standard_decoder.py b/tests/e2e/models/smollm3/test_smollm3_standard_decoder.py new file mode 100644 index 0000000000..4e0a12c4da --- /dev/null +++ b/tests/e2e/models/smollm3/test_smollm3_standard_decoder.py @@ -0,0 +1,358 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for standard_decoder_builder.py — tensor naming contract. + +Builds tiny engines and verifies all I/O tensor names match C++ expectations. +Requires TRT + GPU. + +Trace: ARCH-GRP-001, UD-GRP-DECODER +Intent: Validate standard decoder builder I/O tensor naming contract against C++ runtime expectations +Preconditions: TRT and CUDA GPU are available; synthetic weight dicts match builder requirements +Postconditions: Built engine I/O tensor names exactly match the naming convention expected by C++ runtime +""" + +from __future__ import annotations + +import numpy as np +import pytest + +pytest.importorskip("tensorrt_model_connect", reason="tensorrt_model_connect requires tensorrt") +from tests.builder.conftest import requires_trt + + +def _make_weights(hidden: int, vocab: int, num_layers: int, + attention_size: int, mlp_size: int, + *, mlp_type: str = "swiglu", + position_type: str = "rope", + has_bias: bool = False) -> dict: + """Create a minimal synthetic weight dict for the standard decoder builder.""" + from tensorrt_model_connect.checkpoint_mapper import WeightDict + rng = np.random.RandomState(42) + w = WeightDict() + w["embedding"] = rng.randn(vocab, hidden).astype(np.float32) + + for i in range(num_layers): + p = f"layer.{i}" + w[f"{p}.input_norm"] = rng.randn(hidden).astype(np.float32) + w[f"{p}.post_attn_norm"] = rng.randn(hidden).astype(np.float32) + w[f"{p}.w_q"] = rng.randn(hidden, attention_size).astype(np.float32) + w[f"{p}.w_k"] = rng.randn(hidden, attention_size).astype(np.float32) + w[f"{p}.w_v"] = rng.randn(hidden, attention_size).astype(np.float32) + w[f"{p}.w_o"] = rng.randn(attention_size, hidden).astype(np.float32) + + if mlp_type == "swiglu": + w[f"{p}.w_gate"] = rng.randn(hidden, mlp_size).astype(np.float32) + w[f"{p}.w_up"] = rng.randn(hidden, mlp_size).astype(np.float32) + w[f"{p}.w_down"] = rng.randn(mlp_size, hidden).astype(np.float32) + else: # gelu_fc + w[f"{p}.w_fc1"] = rng.randn(hidden, mlp_size).astype(np.float32) + w[f"{p}.w_fc2"] = rng.randn(mlp_size, hidden).astype(np.float32) + + w["final_norm"] = rng.randn(hidden).astype(np.float32) + w["w_out"] = rng.randn(hidden, vocab).astype(np.float32) + w["_attention_size"] = attention_size + w["_mlp_size"] = mlp_size + + if position_type == "learned": + max_pos = 64 + w["position_embedding"] = rng.randn(max_pos, hidden).astype(np.float32) + + return w + + +def _get_io_names(engine_plan: bytes) -> tuple[list[str], list[str]]: + """Deserialize engine plan and return (input_names, output_names).""" + import tensorrt as trt + logger = trt.Logger(trt.Logger.WARNING) + runtime = trt.Runtime(logger) + engine = runtime.deserialize_cuda_engine(engine_plan) + inputs, outputs = [], [] + for i in range(engine.num_io_tensors): + name = engine.get_tensor_name(i) + mode = engine.get_tensor_mode(name) + if mode == trt.TensorIOMode.INPUT: + inputs.append(name) + else: + outputs.append(name) + return inputs, outputs + + +def _deserialize_engine(engine_plan: bytes): + import tensorrt as trt + logger = trt.Logger(trt.Logger.WARNING) + runtime = trt.Runtime(logger) + return runtime.deserialize_cuda_engine(engine_plan) + + +@requires_trt +class TestTensorNamingContract: + """Verify that built engines have the exact I/O tensor names the C++ runtime expects.""" + + def _build_engine(self, **kwargs): + from tensorrt_model_connect.config import ModelConfig + from tensorrt_model_connect.families.smollm3.standard_decoder_builder import build_standard_decoder_engine + + hidden, vocab, num_layers = 16, 32, 2 + num_heads = 4 + attention_size = hidden + mlp_size = 32 + max_cache = 4 + + config = ModelConfig( + hidden_size=hidden, + vocab_size=vocab, + num_hidden_layers=num_layers, + num_attention_heads=num_heads, + num_key_value_heads=num_heads, + rms_norm_eps=1e-5, + rope_theta=10000.0, + ) + mlp_type = kwargs.get("mlp_type", "swiglu") + position_type = kwargs.get("position_type", "rope") + weights = _make_weights( + hidden, vocab, num_layers, attention_size, mlp_size, + mlp_type=mlp_type, position_type=position_type) + + return build_standard_decoder_engine( + config, weights, max_cache, **kwargs) + + def test_default_rope_swiglu(self): + """Default: RoPE + SwiGLU, standard I/O names.""" + plan = self._build_engine() + inputs, outputs = _get_io_names(plan) + + assert "token_id" in inputs + assert "position_id" in inputs + assert "attention_mask" in inputs + assert "cache_k_0" in inputs + assert "cache_k_1" in inputs + assert "cache_v_0" in inputs + assert "cache_v_1" in inputs + + assert "logits" in outputs + assert "present_k_0" in outputs + assert "present_k_1" in outputs + assert "present_v_0" in outputs + assert "present_v_1" in outputs + + engine = _deserialize_engine(plan) + assert engine.get_tensor_profile_shape("attention_mask", 0) == [ + (1, 5), + (4, 8), + (4, 8), + ] + assert engine.get_tensor_profile_shape("attention_mask", 1) == [ + (1, 5), + (1, 5), + (1, 5), + ] + + def test_layernorm_gelu_fc(self): + """LayerNorm + gelu_fc MLP, same I/O names.""" + plan = self._build_engine( + norm_type="layernorm", mlp_type="gelu_fc", activation="gelu_new") + inputs, outputs = _get_io_names(plan) + + assert "token_id" in inputs + assert "logits" in outputs + assert "present_k_0" in outputs + + def test_learned_positions(self): + """Learned position embeddings, same I/O names.""" + plan = self._build_engine(position_type="learned") + inputs, outputs = _get_io_names(plan) + + assert "token_id" in inputs + assert "position_id" in inputs + assert "logits" in outputs + + def test_alibi_positions(self): + """ALiBi positions, same I/O names.""" + plan = self._build_engine(position_type="alibi") + inputs, outputs = _get_io_names(plan) + + assert "token_id" in inputs + assert "position_id" in inputs + assert "logits" in outputs + + def test_embed_input(self): + """With embed_input=True, extra VL inputs appear.""" + plan = self._build_engine(embed_input=True) + inputs, outputs = _get_io_names(plan) + + assert "input_embed" in inputs + assert "use_input_embed" in inputs + assert "token_id" in inputs + assert "logits" in outputs + + def test_bf16_embed_input_keeps_external_features_fp32(self): + """VL image features stay fp32 while reduced-precision cache uses bf16.""" + import tensorrt as trt + from tensorrt_model_connect.config import ModelConfig + from tensorrt_model_connect.families.smollm3.standard_decoder_builder import ( + build_standard_decoder_engine, + ) + + hidden, vocab, num_layers = 16, 32, 2 + num_heads = 4 + max_cache = 4 + config = ModelConfig( + hidden_size=hidden, + vocab_size=vocab, + num_hidden_layers=num_layers, + num_attention_heads=num_heads, + num_key_value_heads=num_heads, + rms_norm_eps=1e-5, + rope_theta=10000.0, + ) + weights = _make_weights(hidden, vocab, num_layers, hidden, 32) + + plan = build_standard_decoder_engine( + config, weights, max_cache, embed_input=True, precision="bf16") + engine = _deserialize_engine(plan) + + assert engine.get_tensor_dtype("input_embed") == trt.float32 + assert engine.get_tensor_dtype("use_input_embed") == trt.float32 + assert engine.get_tensor_dtype("cache_k_0") == trt.bfloat16 + + def test_debug_layer_outputs(self): + """With debug_layer_outputs=True, per-layer debug outputs appear.""" + plan = self._build_engine(debug_layer_outputs=True) + inputs, outputs = _get_io_names(plan) + + assert "debug_embed" in outputs + assert "debug_hidden_0" in outputs + assert "debug_hidden_1" in outputs + assert "debug_post_attn_0" in outputs + assert "debug_post_attn_1" in outputs + assert "logits" in outputs + + + def test_interleaved_rope(self): + plan = self._build_engine(interleaved_rope=True) + inputs, outputs = _get_io_names(plan) + assert "logits" in outputs + + def test_partial_rotary(self): + plan = self._build_engine(partial_rotary_factor=0.5) + inputs, outputs = _get_io_names(plan) + assert "logits" in outputs + + def test_dynamic_kv_cache_shapes(self): + from tensorrt_model_connect.config import ModelConfig + from tensorrt_model_connect.families.smollm3.standard_decoder_builder import build_standard_decoder_engine + + hidden, vocab, num_layers = 16, 32, 2 + num_heads = 4 + attention_size = hidden + mlp_size = 32 + max_cache = 4 + + config = ModelConfig( + hidden_size=hidden, + vocab_size=vocab, + num_hidden_layers=num_layers, + num_attention_heads=num_heads, + num_key_value_heads=num_heads, + rms_norm_eps=1e-5, + rope_theta=10000.0, + ) + config.raw["dynamic_kv_cache"] = True + weights = _make_weights(hidden, vocab, num_layers, attention_size, mlp_size) + + plan = build_standard_decoder_engine(config, weights, max_cache) + engine = _deserialize_engine(plan) + + assert engine is not None + assert engine.num_optimization_profiles == 2 + assert tuple(engine.get_tensor_shape("attention_mask")) == (-1, -1) + assert tuple(engine.get_tensor_shape("cache_k_0")) == (-1, attention_size) + assert tuple(engine.get_tensor_shape("cache_v_0")) == (-1, attention_size) + assert engine.get_tensor_profile_shape("attention_mask", 0) == [ + (1, 2), + (4, 8), + (4, 8), + ] + assert engine.get_tensor_profile_shape("attention_mask", 1) == [ + (1, 2), + (1, 5), + (1, 5), + ] + + def test_dynamic_kv_cache_multiple_profiles(self): + from tensorrt_model_connect.config import ModelConfig + from tensorrt_model_connect.families.smollm3.standard_decoder_builder import build_standard_decoder_engine + + hidden, vocab, num_layers = 16, 32, 2 + num_heads = 4 + attention_size = hidden + mlp_size = 32 + max_cache = 4 + + config = ModelConfig( + hidden_size=hidden, + vocab_size=vocab, + num_hidden_layers=num_layers, + num_attention_heads=num_heads, + num_key_value_heads=num_heads, + rms_norm_eps=1e-5, + rope_theta=10000.0, + ) + config.raw["dynamic_kv_cache"] = True + config.raw["_dynamic_kv_profile_rows"] = [4, 2, 3] + weights = _make_weights(hidden, vocab, num_layers, attention_size, mlp_size) + + plan = build_standard_decoder_engine(config, weights, max_cache) + engine = _deserialize_engine(plan) + + assert engine is not None + assert engine.num_optimization_profiles == 4 + assert engine.get_tensor_profile_shape("attention_mask", 0) == [ + (1, 2), + (4, 8), + (4, 8), + ] + assert engine.get_tensor_profile_shape("cache_k_0", 0) == [ + (1, attention_size), + (4, attention_size), + (4, attention_size), + ] + for profile, cache_rows in enumerate((2, 3, 4), start=1): + assert engine.get_tensor_profile_shape( + "attention_mask", profile + ) == [ + (1, 2), + (1, cache_rows + 1), + (1, cache_rows + 1), + ] + assert engine.get_tensor_profile_shape("cache_k_0", profile) == [ + (1, attention_size), + (cache_rows, attention_size), + (cache_rows, attention_size), + ] + + def test_dynamic_kv_cache_rejects_alibi(self): + from tensorrt_model_connect.config import ModelConfig + from tensorrt_model_connect.families.smollm3.standard_decoder_builder import build_standard_decoder_engine + + hidden, vocab, num_layers = 16, 32, 2 + num_heads = 4 + attention_size = hidden + mlp_size = 32 + max_cache = 4 + + config = ModelConfig( + hidden_size=hidden, + vocab_size=vocab, + num_hidden_layers=num_layers, + num_attention_heads=num_heads, + num_key_value_heads=num_heads, + rms_norm_eps=1e-5, + rope_theta=10000.0, + ) + config.raw["dynamic_kv_cache"] = True + weights = _make_weights(hidden, vocab, num_layers, attention_size, mlp_size) + + with pytest.raises(ValueError, match="ALiBi"): + build_standard_decoder_engine(config, weights, max_cache, position_type="alibi") diff --git a/tests/e2e/models/smollm3/thresholds/smollm3-3b.json b/tests/e2e/models/smollm3/thresholds/smollm3-3b.json new file mode 100644 index 0000000000..6793c27631 --- /dev/null +++ b/tests/e2e/models/smollm3/thresholds/smollm3-3b.json @@ -0,0 +1,13 @@ +{ + "threshold_overrides": { + "layer_atol": 0.05, + "logit_atol": 0.001, + "logit_cosine_p5": 0.99, + "logit_rel_l2_p95": 0.05, + "normalized_text_edit_distance": 0.2, + "stable_margin": 0.1, + "stable_top1_match_rate": 0.9, + "token_agreement_rate": 0.8, + "unstable_topk_hit_rate": 0.8 + } +} diff --git a/tests/tools/test_perf_matrix.py b/tests/tools/test_perf_matrix.py index 4e901ad254..2523c15540 100644 --- a/tests/tools/test_perf_matrix.py +++ b/tests/tools/test_perf_matrix.py @@ -51,6 +51,10 @@ def _suite_for_cases(cases, *, exclusions=None): "Pinned BF16 build and Hugging Face parity qualification are present, but " "no matching release-performance workload or receipt has been collected." ) +SMOLLM3_EXCLUSION_REASON = ( + "Dense SmolLM3 functional and reference-parity qualification is present, but " + "this change does not add a matching release-performance workload or receipt." +) TASK_ADAPTERS = { "bark.generate_audio": "hf-transformers-tts", "bert.embed": "hf-transformers-embedding", @@ -287,6 +291,7 @@ def test_release_suite_covers_every_non_l0_ready_model_profile() -> None: "lfm2-700m": LFM2_EXCLUSION_REASON, "k2-horizon-7b": K2_HORIZON_EXCLUSION_REASON, "minimax-h3-768p": MINIMAX_H3_EXCLUSION_REASON, + "smollm3-3b": SMOLLM3_EXCLUSION_REASON, } assert all( set(entry["workload"]) <= {"testcase", "request", "runtime"} for entry in raw_entries @@ -2220,8 +2225,8 @@ def preflight_after_pending_report(cases, options): expected_catalog_coverage = { "total_profiles": len(catalog_entries), "ready_profiles": catalog_counts["ready"], - "release_profiles": catalog_counts["ready"] - excluded_l0_profiles - 7, - "explicitly_excluded_profiles": 7, + "release_profiles": catalog_counts["ready"] - excluded_l0_profiles - 8, + "explicitly_excluded_profiles": 8, "explicit_exclusions": [ { "model": "lfm2-1.2b", @@ -2251,6 +2256,10 @@ def preflight_after_pending_report(cases, options): "model": "minimax-h3-768p", "reason": MINIMAX_H3_EXCLUSION_REASON, }, + { + "model": "smollm3-3b", + "reason": SMOLLM3_EXCLUSION_REASON, + }, ], "excluded_l0_profiles": excluded_l0_profiles, "distributed_profiles": catalog_counts["distributed"], diff --git a/tests/validation/model_workloads.yaml b/tests/validation/model_workloads.yaml index 0d98cd50ba..950a580ff2 100644 --- a/tests/validation/model_workloads.yaml +++ b/tests/validation/model_workloads.yaml @@ -285,6 +285,8 @@ models: workloads: [sana_wm_benchmark_diffusion_video] segformer-b0-ade: workloads: [ade20k_semantic_segmentation] + smollm3-3b: + workloads: [mmlu_continuation_parity] stablelm2-1.6b: workloads: [mmlu_continuation_parity] starcoder2-3b: diff --git a/website/data/hf-model-metadata.json b/website/data/hf-model-metadata.json index d7edf1413a..4b87838a1f 100644 --- a/website/data/hf-model-metadata.json +++ b/website/data/hf-model-metadata.json @@ -91,6 +91,17 @@ ], "architecture_source": "config.architectures" }, + { + "hf_id": "HuggingFaceTB/SmolLM3-3B", + "revision": "a07cc9a04f16550a088caea529712d1d335b0ac1", + "revision_source": "declared", + "metadata_file": "config.json", + "model_type": "smollm3", + "architectures": [ + "SmolLM3ForCausalLM" + ], + "architecture_source": "config.architectures" + }, { "hf_id": "IFM/K2-Horizon-7B", "revision": "586b03f0fd1fbbf2f13eeafc33749e95ae34dd10",