diff --git a/benchmarks/performance/release.yaml b/benchmarks/performance/release.yaml index 21ef60ef5..e36d6f634 100644 --- a/benchmarks/performance/release.yaml +++ b/benchmarks/performance/release.yaml @@ -758,6 +758,16 @@ entries: runner: hf-transformers mode: torch-compile compile_scope: model.forward + - id: qwen3_8.generate + family: qwen3_8 + operation: generate + model: qwen38-27b + workload: + testcase: qwen38-27b + baseline: + runner: hf-transformers + mode: torch-compile + compile_scope: model.forward - id: qwen3_omni.generate_audio family: qwen3_omni operation: generate_audio diff --git a/python/tensorrt_model_connect/families/qwen3_8/MODEL.toml b/python/tensorrt_model_connect/families/qwen3_8/MODEL.toml new file mode 100644 index 000000000..89427490a --- /dev/null +++ b/python/tensorrt_model_connect/families/qwen3_8/MODEL.toml @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +id = "qwen3_8" +plugin = "qwen3_8" +module = "plugin" +aliases = [ + "qwen38", + "qwen3.8", + "qwen3_8", +] +prefixes = [ + "qwen3_8", +] +# Qwen3.8 checkpoints reuse the Qwen3.5 `model_type`/`architectures` strings, so +# alias matching alone cannot separate the two families. Declaring the shared +# architecture pattern routes config-shaped lookups through this family first; +# `Qwen38Plugin.matches_config` then claims only genuine Qwen3.8 checkpoints and +# leaves Qwen3.5 to fall through to its own family. +architecture_patterns = [ + "qwen3_5forconditionalgeneration", +] +debug_runner = "debug_runner.py|runner_from_bundle" +debug_runtime_strategies = [ + "qwen3_8_hybrid_mamba_attention", +] diff --git a/python/tensorrt_model_connect/families/qwen3_8/__init__.py b/python/tensorrt_model_connect/families/qwen3_8/__init__.py new file mode 100644 index 000000000..136100f92 --- /dev/null +++ b/python/tensorrt_model_connect/families/qwen3_8/__init__.py @@ -0,0 +1,62 @@ +# 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) + 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/qwen3_8/checkpoint_mapper.py b/python/tensorrt_model_connect/families/qwen3_8/checkpoint_mapper.py new file mode 100644 index 000000000..e6d1daebd --- /dev/null +++ b/python/tensorrt_model_connect/families/qwen3_8/checkpoint_mapper.py @@ -0,0 +1,199 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""1:1 port of standard_checkpoint_mapper.cpp + tensor_math.cpp to Python. + +Loads HF safetensors and maps keys to the flat weight dict expected by +standard_decoder_builder.py. All projections are transposed from HF +[out, in] layout to [in, out] for TRT matmul. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np + +# Register bfloat16 dtype with numpy (needed for safetensors without torch). +try: + import ml_dtypes # noqa: F401 +except ImportError: + pass + +from safetensors import safe_open + + +def _target_np_dtype(precision: str) -> np.dtype: + """Map precision string to numpy dtype for weight storage.""" + if precision in ("fp16", "bf16"): + return np.float16 + return np.float32 + + +def _transpose_2d(arr: np.ndarray, name: str, precision: str = "fp32") -> np.ndarray: + """Transpose [rows, cols] -> [cols, rows] in C-contiguous target dtype.""" + if arr.ndim != 2: + raise ValueError(f"Expected rank-2 tensor for transpose: {name}") + return np.ascontiguousarray(arr.T, dtype=_target_np_dtype(precision)) + + +class WeightDict(dict): + """A dict mapping logical weight names to flat float32 arrays. + + Keys follow the convention used by standard_decoder_builder.py: + - embedding: [vocab, hidden] + - layer.{i}.input_norm: [hidden] + - layer.{i}.w_q: [hidden, attention_size] + - layer.{i}.w_k: [hidden, kv_attention_size] + - layer.{i}.w_v: [hidden, kv_attention_size] + - layer.{i}.q_bias: [attention_size] (optional) + - layer.{i}.k_bias: [kv_attention_size] (optional) + - layer.{i}.v_bias: [kv_attention_size] (optional) + - layer.{i}.q_norm: [attention_size] (optional) + - layer.{i}.k_norm: [kv_attention_size] (optional) + - layer.{i}.w_o: [attention_size, hidden] + - layer.{i}.post_attn_norm: [hidden] + - layer.{i}.w_gate: [hidden, mlp_size] + - layer.{i}.w_up: [hidden, mlp_size] + - layer.{i}.w_down: [mlp_size, hidden] + - final_norm: [hidden] + - w_out: [hidden, vocab] + """ + + +# --------------------------------------------------------------------------- +# Safetensors I/O helpers +# --------------------------------------------------------------------------- + + +def _detect_framework() -> str: + """Use 'torch' if available (handles BF16 natively), else 'numpy'.""" + try: + import torch # noqa: F401 + + return "torch" + except ImportError: + return "numpy" + + +class _TorchBinReader: + """Adapter that wraps a pytorch .bin state dict with the safetensors reader + interface (keys() / get_tensor()).""" + + def __init__(self, path: Path): + import torch + + self._state = torch.load(str(path), map_location="cpu", weights_only=True) + + def keys(self) -> list[str]: + return list(self._state.keys()) + + def get_tensor(self, name: str): + return self._state[name] + + +class _ReaderCollection(list): + """Reader list with a cached tensor-name -> reader lookup table.""" + + def __init__(self, readers: list, *, tensor_map: dict[str, object] | None = None): + super().__init__(readers) + if tensor_map is None: + tensor_map = {} + for reader in readers: + for key in reader.keys(): + tensor_map[key] = reader + self.tensor_map = tensor_map + + +def _open_safetensors(model_dir: Path) -> list: + """Open all safetensor shards (or pytorch .bin) in a model directory.""" + fw = _detect_framework() + single = model_dir / "model.safetensors" + if single.exists(): + return _ReaderCollection([safe_open(str(single), framework=fw)]) + + index_path = model_dir / "model.safetensors.index.json" + if index_path.exists(): + import json + + index = json.loads(index_path.read_text()) + weight_map = index.get("weight_map", {}) + shard_files = sorted(set(weight_map.values())) + readers_by_file = { + shard: safe_open(str(model_dir / shard), framework=fw) for shard in shard_files + } + tensor_map = {name: readers_by_file[shard] for name, shard in weight_map.items()} + return _ReaderCollection( + [readers_by_file[shard] for shard in shard_files], + tensor_map=tensor_map, + ) + + # Diffusers format: diffusion_pytorch_model.safetensors + diff_single = model_dir / "diffusion_pytorch_model.safetensors" + if diff_single.exists(): + return _ReaderCollection([safe_open(str(diff_single), framework=fw)]) + + diff_index = model_dir / "diffusion_pytorch_model.safetensors.index.json" + if diff_index.exists(): + import json + + index = json.loads(diff_index.read_text()) + weight_map = index.get("weight_map", {}) + shard_files = sorted(set(weight_map.values())) + readers_by_file = { + shard: safe_open(str(model_dir / shard), framework=fw) for shard in shard_files + } + tensor_map = {name: readers_by_file[shard] for name, shard in weight_map.items()} + return _ReaderCollection( + [readers_by_file[shard] for shard in shard_files], + tensor_map=tensor_map, + ) + + # Fallback: pytorch_model.bin (older HF models) + bin_single = model_dir / "pytorch_model.bin" + if bin_single.exists(): + return _ReaderCollection([_TorchBinReader(bin_single)]) + + raise FileNotFoundError( + f"No model.safetensors, index.json, or pytorch_model.bin in {model_dir}" + ) + + +def _has_tensor(readers: list, name: str) -> bool: + tensor_map = getattr(readers, "tensor_map", None) + if tensor_map is not None: + return name in tensor_map + for r in readers: + if name in r.keys(): + return True + return False + + +def _to_numpy_fp32(t) -> np.ndarray: + """Convert a safetensors/torch tensor to numpy float32 with minimal copies.""" + if hasattr(t, "numpy"): + dtype = getattr(t, "dtype", None) + if str(dtype) == "torch.float32": + return t.numpy() + return t.float().numpy() + + dtype_str = str(t.dtype) + if t.dtype == np.uint16 or dtype_str == "bfloat16": + t = t.view(np.uint16).astype(np.uint32) << 16 + return t.view(np.float32) + if dtype_str == "float16": + return t.astype(np.float32) + return np.asarray(t, dtype=np.float32) + + +def _load_tensor(readers: list, name: str) -> np.ndarray: + tensor_map = getattr(readers, "tensor_map", None) + if tensor_map is not None: + reader = tensor_map.get(name) + if reader is None: + raise KeyError(f"Tensor not found: {name}") + return _to_numpy_fp32(reader.get_tensor(name)) + for r in readers: + if name in r.keys(): + return _to_numpy_fp32(r.get_tensor(name)) + raise KeyError(f"Tensor not found: {name}") diff --git a/python/tensorrt_model_connect/families/qwen3_8/config.py b/python/tensorrt_model_connect/families/qwen3_8/config.py new file mode 100644 index 000000000..9ba70ea29 --- /dev/null +++ b/python/tensorrt_model_connect/families/qwen3_8/config.py @@ -0,0 +1,232 @@ +# 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 _token_id(value: object) -> int: + """Normalize a config token id to an int, using -1 for "unset". + + A plain ``value or -1`` would map a declared id of ``0`` to ``-1``, because + ``0`` is falsy, silently dropping EOS detection or padding for checkpoints + that use token 0. Lists (multiple stop ids) take the first entry, matching + the runtime's own first-of-array behavior. + """ + if isinstance(value, list): + value = value[0] if value else None + if isinstance(value, bool) or not isinstance(value, int): + return -1 + return value + + +@dataclass +class ModelConfig: + """Parsed model architecture from HF config.json.""" + + model_type: str = "" + architectures: list[str] = field(default_factory=list) + vocab_size: int = 0 + hidden_size: int = 0 + intermediate_size: int = 0 + num_hidden_layers: int = 0 + num_attention_heads: int = 1 + num_key_value_heads: int = 1 + rms_norm_eps: float = 1e-5 + rope_theta: float = 10000.0 + bos_token_id: int = -1 + eos_token_id: int = -1 + pad_token_id: int = -1 + tie_word_embeddings: bool = False + max_position_embeddings: int = 8192 + hidden_act: str = "" + + # Explicit head_dim from config.json (0 = not set, fall back to computed). + _head_dim: int = 0 + + # Raw JSON dict for family-specific fields + raw: dict = field(default_factory=dict, repr=False) + + @property + def head_dim(self) -> int: + if self._head_dim > 0: + return self._head_dim + if self.num_attention_heads <= 0: + return 0 + return self.hidden_size // self.num_attention_heads + + @property + def attention_size(self) -> int: + return self.num_attention_heads * self.head_dim + + @staticmethod + def from_json(text: str) -> ModelConfig: + d = json.loads(text) + + # Some multimodal configs nest decoder fields under "text_config". + # Merge text_config into top level so standard key lookup works. + # Preserve top-level model_type and architectures (these identify the + # top-level model, not the nested decoder). + original_raw = d + text_config = d.get("text_config") + if text_config and isinstance(text_config, dict): + top_model_type = d.get("model_type") + top_architectures = d.get("architectures") + merged = {**d, **text_config} + if top_model_type: + merged["model_type"] = top_model_type + if top_architectures: + merged["architectures"] = top_architectures + d = merged + + # Some multimodal configs nest the language decoder config under + # "language_config". Merge into top level like text_config. + if not d.get("hidden_size"): + lang_config = d.get("language_config") + if isinstance(lang_config, dict): + top_model_type = d.get("model_type") + top_architectures = d.get("architectures") + top_vision_config = d.get("vision_config") + merged = {**d, **lang_config} + if top_model_type: + merged["model_type"] = top_model_type + if top_architectures: + merged["architectures"] = top_architectures + if top_vision_config: + merged["vision_config"] = top_vision_config + d = merged + + # Some multimodal configs nest LLM config under "llm_config". + # Merge into top level like text_config, preserving top-level + # model_type, architectures, and vision_config. + if not d.get("hidden_size"): + llm_config = d.get("llm_config") + if isinstance(llm_config, dict): + top_model_type = d.get("model_type") + top_architectures = d.get("architectures") + top_vision_config = d.get("vision_config") + merged = {**d, **llm_config} + if top_model_type: + merged["model_type"] = top_model_type + if top_architectures: + merged["architectures"] = top_architectures + if top_vision_config: + merged["vision_config"] = top_vision_config + d = merged + + # Some multimodal audio/text configs nest the primary decoder config + # under thinker_config.text_config. If top-level hidden_size is + # still missing after the text_config merge above, look there. + if not d.get("hidden_size"): + thinker_cfg = d.get("thinker_config") + if isinstance(thinker_cfg, dict): + thinker_text = thinker_cfg.get("text_config") + if isinstance(thinker_text, dict): + top_model_type = d.get("model_type") + top_architectures = d.get("architectures") + merged = {**d, **thinker_text} + if top_model_type: + merged["model_type"] = top_model_type + if top_architectures: + merged["architectures"] = top_architectures + # Also propagate vision_config from thinker_config + # so VL pipelines can find it. + if "vision_config" not in merged and "vision_config" in thinker_cfg: + merged["vision_config"] = thinker_cfg["vision_config"] + d = merged + + # Handle non-standard config key names: + # GPT-2: n_embd, n_head, n_layer, n_inner + # XGLM/Bloom: d_model, attention_heads, num_layers, ffn_dim + # DistilBERT: dim, n_heads, n_layers, hidden_dim + hidden_size = (d.get("hidden_size", 0) or d.get("n_embd", 0) + or d.get("d_model", 0) or d.get("n_embed", 0) + or d.get("dim", 0)) + num_heads = (d.get("num_attention_heads", 0) or d.get("n_head", 0) + or d.get("attention_heads", 0) or d.get("num_heads", 0) + or d.get("n_heads", 0) or d.get("decoder_attention_heads", 0) or 1) + num_layers = (d.get("num_hidden_layers", 0) or d.get("n_layer", 0) + or d.get("num_layers", 0) or d.get("n_layers", 0)) + intermediate = (d.get("intermediate_size", 0) + or d.get("n_inner", 0) + or d.get("ffn_dim", 0) + or d.get("hidden_dim", 0) + or hidden_size * 4) + + # Norm epsilon: try rms_norm_eps, then layer_norm_epsilon, then + # layer_norm_eps, then norm_epsilon, then norm_eps. + eps = (d.get("rms_norm_eps") + or d.get("layer_norm_epsilon") + or d.get("layer_norm_eps") + or d.get("norm_epsilon") + or d.get("norm_eps") + or 1e-5) + + # rope_theta: check top-level first, then rope_parameters dict + # (some model configs store it there), + # then rope_scaling dict. + rope_theta = d.get("rope_theta", None) + if rope_theta is None: + rope_params = d.get("rope_parameters") + if isinstance(rope_params, dict): + rope_theta = rope_params.get("rope_theta", 10000.0) + else: + rope_scaling = d.get("rope_scaling") + if isinstance(rope_scaling, dict): + rope_theta = rope_scaling.get("rope_theta", 10000.0) + else: + rope_theta = 10000.0 + rope_theta = float(rope_theta) + + architecture = d.get("architecture", "") + architectures = d.get("architectures", []) + if not architectures and architecture: + architectures = [architecture] + + return ModelConfig( + model_type=d.get("model_type", "") or architecture, + architectures=architectures, + vocab_size=d.get("vocab_size", 0), + hidden_size=hidden_size or d.get("num_features", 0), + intermediate_size=intermediate, + num_hidden_layers=num_layers, + num_attention_heads=num_heads, + num_key_value_heads=d.get("num_key_value_heads", num_heads), + rms_norm_eps=eps, + rope_theta=rope_theta, + bos_token_id=_token_id(d.get("bos_token_id")), + eos_token_id=_token_id(d.get("eos_token_id")), + pad_token_id=_token_id(d.get("pad_token_id")), + 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/qwen3_8/cpu_profile_matrix.py b/python/tensorrt_model_connect/families/qwen3_8/cpu_profile_matrix.py new file mode 100644 index 000000000..cceb7e006 --- /dev/null +++ b/python/tensorrt_model_connect/families/qwen3_8/cpu_profile_matrix.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU profile matrix defaults owned by the Qwen3.8 family.""" + +from __future__ import annotations + + +def cpu_profile_matrix_specs() -> list[dict]: + return [{ + "order": 51, + "strategy": "qwen3_8_hybrid_mamba_attention", + "label": "hybrid_mamba_attn\n(qwen38-27b)", + "hf_id": "Qwen/Qwen3.8-27B", + "bundle": "qwen38-27b.bundle", + "runner": "decoder", + }] diff --git a/python/tensorrt_model_connect/families/qwen3_8/debug_runner.py b/python/tensorrt_model_connect/families/qwen3_8/debug_runner.py new file mode 100644 index 000000000..c510cf160 --- /dev/null +++ b/python/tensorrt_model_connect/families/qwen3_8/debug_runner.py @@ -0,0 +1,479 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Qwen3.8-owned hybrid debug runner implementation.""" + +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] + +# Mirrors kMaskedScore in src/runtime/models/qwen3_8/kv_cache.cpp. +_MASKED_SCORE = -1.0e4 + + +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 _check_cuda(status): + if cudart is None: + raise RuntimeError("cuda-python is required for 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 _require_trt_runtime() -> None: + if trt is None: + raise ImportError("tensorrt is required for debug_runner execution") + if cudart is None: + raise ImportError("cuda-python is required for debug_runner execution") + + +def load_engine_from_bundle( + bundle_path: str, + section_name: str = "engine_plan", +) -> tuple[bytes, dict]: + """Load this family's engine plan bytes and bundle metadata.""" + import json + import struct + + with open(bundle_path, "rb") as f: + magic = f.read(8) + if magic != b"BUNDLE\x01\x00": + raise ValueError(f"Not a valid .bundle artifact: {bundle_path}") + header_len = struct.unpack(" bytes | None: + """Load a named raw section from this family's .bundle artifact.""" + import json + import struct + + with open(bundle_path, "rb") as f: + magic = f.read(8) + if magic != b"BUNDLE\x01\x00": + raise ValueError(f"Not a valid .bundle artifact: {bundle_path}") + header_len = struct.unpack(" dict: + """Load and parse this family's config.json from a .bundle artifact.""" + import json + + data = load_section_from_bundle(bundle_path, "config.json") + if data is None: + return {} + return json.loads(data.decode("utf-8")) + + +class HybridTrtRunner: + """Device-resident hybrid TRT inference runner for models with mixed + recurrent (DeltaNet/Mamba) + attention layers. + + Combines recurrent conv/SSM state management with KV cache + position + tracking for hybrid recurrent-attention models. + """ + + def __init__( + self, + engine_plan: bytes, + max_cache_length: int, + num_mamba_layers: int, + num_attention_layers: int, + distributed_communicator: object | None = None, + ): + _require_trt_runtime() + self.max_cache_length = max_cache_length + self.num_mamba_layers = num_mamba_layers + self.num_attention_layers = num_attention_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") + + # Auto-detect state dimensions from engine tensor shapes + if num_mamba_layers > 0: + conv_shape = tuple(self.engine.get_tensor_shape("conv_state_0")) + ssm_shape = tuple(self.engine.get_tensor_shape("ssm_state_0")) + else: + conv_shape = (0,) + ssm_shape = (0,) + + if num_attention_layers > 0: + cache_shape = tuple(self.engine.get_tensor_shape("cache_k_0")) + self.attention_size = cache_shape[1] + cache_dtype = self.engine.get_tensor_dtype("cache_k_0") + self._cache_elem_bytes = _trt_itemsize(cache_dtype) + else: + self.attention_size = 0 + self._cache_elem_bytes = 4 + + err, self.stream = cudart.cudaStreamCreate() + _check_cuda(err) + + self.cache_length = 0 + attention_window = max_cache_length + 1 + + # Discover debug output tensor names + 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 + if (name != "logits" + and not name.startswith("present_conv_") + and not name.startswith("present_ssm_") + and not name.startswith("present_k_") + and not name.startswith("present_v_")): + self._debug_output_names.append(name) + + # --- Mamba/DeltaNet state buffers --- + self._conv_state_bytes = int(np.prod(conv_shape)) * 4 if num_mamba_layers > 0 else 0 + self._ssm_state_bytes = int(np.prod(ssm_shape)) * 4 if num_mamba_layers > 0 else 0 + + self._d_conv_state: list[int] = [] + self._d_ssm_state: list[int] = [] + self._d_present_conv: list[int] = [] + self._d_present_ssm: list[int] = [] + for _ in range(num_mamba_layers): + for lst, sz in [(self._d_conv_state, self._conv_state_bytes), + (self._d_ssm_state, self._ssm_state_bytes), + (self._d_present_conv, self._conv_state_bytes), + (self._d_present_ssm, self._ssm_state_bytes)]: + err, ptr = cudart.cudaMalloc(sz) + _check_cuda(err) + lst.append(ptr) + + # --- KV cache buffers --- + row_bytes = self.attention_size * self._cache_elem_bytes + cache_bytes = max_cache_length * row_bytes + + self._d_cache_k: list[int] = [] + self._d_cache_v: list[int] = [] + self._d_present_k: list[int] = [] + self._d_present_v: list[int] = [] + for _ in range(num_attention_layers): + for lst, sz in [(self._d_cache_k, cache_bytes), + (self._d_cache_v, cache_bytes), + (self._d_present_k, row_bytes), + (self._d_present_v, row_bytes)]: + err, ptr = cudart.cudaMalloc(sz) + _check_cuda(err) + lst.append(ptr) + + # --- Small I/O --- + 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) + + self._h_mask = np.zeros((1, attention_window), dtype=np.float32) + err, self._d_mask = cudart.cudaMalloc(attention_window * 4) + _check_cuda(err) + + 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) + + # Debug output 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 all state + for i in range(num_mamba_layers): + _check_cuda(cudart.cudaMemsetAsync( + self._d_conv_state[i], 0, self._conv_state_bytes, self.stream)[0]) + _check_cuda(cudart.cudaMemsetAsync( + self._d_ssm_state[i], 0, self._ssm_state_bytes, self.stream)[0]) + for i in range(num_attention_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) + + def step(self, token_id: int) -> dict[str, np.ndarray]: + """Run one hybrid decode step.""" + 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). + # Must equal kMaskedScore in src/runtime/models/qwen3_8/kv_cache.cpp. + # -1e9 would differ from the native runtime, and is outside the FP16 + # range: an FP16 engine turns it into -inf, which yields NaN after a + # softmax over a fully masked row. + position_id = min(self.cache_length, self.max_cache_length) + self._h_mask[:] = _MASKED_SCORE + valid = min(self.cache_length, self.max_cache_length) + self._h_mask[0, :valid] = 0.0 + self._h_mask[0, -1] = 0.0 + + self._h_token_id[0] = token_id + self._h_position_id[0] = position_id + + # H2D: small inputs + 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) + + # 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) + + for i in range(self.num_mamba_layers): + self.context.set_tensor_address( + f"conv_state_{i}", self._d_conv_state[i]) + self.context.set_tensor_address( + f"ssm_state_{i}", self._d_ssm_state[i]) + self.context.set_tensor_address( + f"present_conv_{i}", self._d_present_conv[i]) + self.context.set_tensor_address( + f"present_ssm_{i}", self._d_present_ssm[i]) + + for i in range(self.num_attention_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]) + + # Execute + self.context.execute_async_v3(stream) + + # D2D state update: conv/ssm (direct replacement) + for i in range(self.num_mamba_layers): + cudart.cudaMemcpyAsync( + self._d_conv_state[i], self._d_present_conv[i], + self._conv_state_bytes, D2D, stream) + cudart.cudaMemcpyAsync( + self._d_ssm_state[i], self._d_present_ssm[i], + self._ssm_state_bytes, D2D, stream) + + # D2D cache update: KV (append or shift) + row_bytes = self.attention_size * self._cache_elem_bytes + for i in range(self.num_attention_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: + # cudaMemcpyAsync is undefined for overlapping ranges, so the + # cache-full shift stages through scratch, matching + # Qwen38KvCache::advance() in the native runtime. One buffer + # serves every layer because the copies are stream-ordered. + shift_bytes = (self.max_cache_length - 1) * row_bytes + if shift_bytes > 0: + scratch = self._shift_scratch(shift_bytes) + cudart.cudaMemcpyAsync( + scratch, cache_buf + row_bytes, shift_bytes, D2D, stream) + cudart.cudaMemcpyAsync( + cache_buf, scratch, shift_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) + + 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 state buffers and reset cache_length.""" + cache_bytes = self.max_cache_length * self.attention_size * self._cache_elem_bytes + for i in range(self.num_mamba_layers): + _check_cuda(cudart.cudaMemsetAsync( + self._d_conv_state[i], 0, self._conv_state_bytes, self.stream)[0]) + _check_cuda(cudart.cudaMemsetAsync( + self._d_ssm_state[i], 0, self._ssm_state_bytes, self.stream)[0]) + for i in range(self.num_attention_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.""" + all_results = [] + for tid in input_ids: + all_results.append(self.step(tid)) + for _ in range(max_new_tokens): + next_token = int(np.argmax(all_results[-1]["logits"].flatten())) + all_results.append(self.step(next_token)) + return all_results + + def _shift_scratch(self, nbytes: int) -> int: + """Device scratch for the cache-full shift, allocated on first overflow. + + Mirrors Qwen38KvCache::shift_scratch() in the native runtime: a cache + that never fills never pays for the allocation. + """ + current = getattr(self, "_d_shift_scratch", 0) + if current and getattr(self, "_shift_scratch_bytes", 0) >= nbytes: + return current + if current: + cudart.cudaFree(current) + err, ptr = cudart.cudaMalloc(nbytes) + _check_cuda(err) + self._d_shift_scratch = ptr + self._shift_scratch_bytes = nbytes + return ptr + + 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_conv_state) + bufs.extend(self._d_ssm_state) + bufs.extend(self._d_present_conv) + bufs.extend(self._d_present_ssm) + bufs.extend(self._d_cache_k) + bufs.extend(self._d_cache_v) + bufs.extend(self._d_present_k) + bufs.extend(self._d_present_v) + for d_ptr in self._d_debug.values(): + bufs.append(d_ptr) + if getattr(self, "_d_shift_scratch", 0): + bufs.append(self._d_shift_scratch) + 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, +) -> HybridTrtRunner: + del runtime_strategy, bundle_path + return HybridTrtRunner( + engine_plan=engine_plan, + max_cache_length=header["max_cache_length"], + num_mamba_layers=config.get("num_mamba_layers", 0), + num_attention_layers=config.get("num_attention_layers", 0), + distributed_communicator=distributed_communicator, + ) diff --git a/python/tensorrt_model_connect/families/qwen3_8/graph_blocks.py b/python/tensorrt_model_connect/families/qwen3_8/graph_blocks.py new file mode 100644 index 000000000..b4d19b961 --- /dev/null +++ b/python/tensorrt_model_connect/families/qwen3_8/graph_blocks.py @@ -0,0 +1,149 @@ +# 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}, 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_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 diff --git a/python/tensorrt_model_connect/families/qwen3_8/graph_ops.py b/python/tensorrt_model_connect/families/qwen3_8/graph_ops.py new file mode 100644 index 000000000..ff7248a89 --- /dev/null +++ b/python/tensorrt_model_connect/families/qwen3_8/graph_ops.py @@ -0,0 +1,877 @@ +# 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 + +import numpy as np +from tensorrt_model_connect import trt_compat + + +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 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, +) -> 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) + 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_l2_norm( + network: trt.INetworkDefinition, + inp: trt.ITensor, + reduce_axis: int, + eps: float = 1e-12, + dtype: np.dtype = np.float32, +) -> trt.ITensor: + """L2 normalize: x / max(||x||_2, eps) along reduce_axis. + + Used for DeltaNet Q/K normalization (Gated DeltaNet architecture). + + 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) + sq = network.add_elementwise(inp, inp, trt.ElementWiseOperation.PROD) + sum_sq = network.add_reduce( + sq.get_output(0), trt.ReduceOperation.SUM, 1 << reduce_axis, keep_dims=True + ) + norm = network.add_unary(sum_sq.get_output(0), trt.UnaryOperation.SQRT) + # max(norm, eps) to avoid division by zero + eps_const = add_constant( + network, (1,) * (reduce_axis + 1), np.array([eps], dtype=np.float32), dtype=np.float32 + ) + safe_norm = network.add_elementwise(norm.get_output(0), eps_const, trt.ElementWiseOperation.MAX) + recip = network.add_unary(safe_norm.get_output(0), trt.UnaryOperation.RECIP) + normalized = network.add_elementwise(inp, recip.get_output(0), trt.ElementWiseOperation.PROD) + result = normalized.get_output(0) + if need_cast: + result = _cast_back_to_trt_dtype(network, result, output_dtype) + return result + + +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}") + + +# 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_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, +) -> 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). + + 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) + table = np.full((max_cache_length, half), default, dtype=np.float32) + for pos in range(max_cache_length): + for d in range(half): + # For both interleaved and rotate-half the frequency index is d + # (the distinction only affects which input pair is rotated; the + # freq assignment per half-dim is the same). + exponent = (2.0 * d) / rotary_ndims + inv_freq = rope_theta ** (-exponent) + angle = pos * inv_freq + table[pos, d] = np.cos(angle) if cosine else np.sin(angle) + return table + + +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_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, + 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 (IRotaryEmbeddingLayer with position_ids): + 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 + 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. + 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) + + # 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 = network.add_rotary_embedding( + inp_4d, + cos_cache_2d, + sin_cache_2d, + interleaved, + rotary_embedding_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 _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 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 diff --git a/python/tensorrt_model_connect/families/qwen3_8/plugin.py b/python/tensorrt_model_connect/families/qwen3_8/plugin.py new file mode 100644 index 000000000..b7ebfe1a4 --- /dev/null +++ b/python/tensorrt_model_connect/families/qwen3_8/plugin.py @@ -0,0 +1,1344 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Qwen3.8 family plugin -- Hybrid Gated DeltaNet + self-attention decoder. + +Qwen3.8 is an architectural re-release of Qwen3.5: the checkpoints keep +`model_type: "qwen3_5"` and `architectures: ["Qwen3_5ForConditionalGeneration"]`, +and the tensor layout is unchanged. It is a separate family here because the +model family is this project's unit of ownership (see AGENTS.md) -- Qwen3.8 must +be implementable, validatable and revertable without touching qwen3_5. + +Family dispatch therefore cannot key on `model_type`; see +`Qwen38Plugin.matches_config` for the config-body discriminator. + +Qwen3.8 uses a heterogeneous layer stack with two layer types defined by +text_config.layer_types (list of strings): + "linear_attention" = Gated DeltaNet layer (linear attention with delta rule) + "full_attention" = Standard self-attention layer (GQA, partial RoPE, output gating) + +Qwen3.8-27B: 64 layers (48 DeltaNet + 16 self-attention), hidden 5120, +24 query heads / 4 KV heads, head_dim 256, vocab 248320, untied lm_head. +Full-attention layers appear every 4th layer (indices 3, 7, 11, ...), i.e. +`full_attention_interval: 4`. + +DeltaNet head expansion: `linear_num_key_heads: 16` carries Q/K while +`linear_num_value_heads: 48` carries V, so Q/K are broadcast 3x to meet V. +(Qwen3.5-9B already runs this path at 2x and Qwen3.5-2B at 1x; only the ratio +differs here.) + +Config keys Qwen3.8 adds that this graph deliberately ignores: + - `output_gate_type: "swish"` -- inert. transformers v5.8.0 (the version the + checkpoint declares) has no such field in Qwen3_5Config; the reference + gates the DeltaNet norm with `config.hidden_act` ("silu", == swish) and the + attention output with `sigmoid`, which is exactly what this graph does. + - `mtp_num_hidden_layers: 1` -- the `mtp.*` speculative-decoding head is + present in the checkpoint and is not part of the decoder graph. + - vision tower (`model.visual.*`) and mrope image/video sections -- the + text-only decoder path does not consume them. + +Key architecture details: + + DeltaNet layers: + - in_proj_qkv -> conv1d step -> SiLU -> split Q[nkv x dim], K[nkv x dim], V[nheads x dim] + - L2-norm Q and K + - keep compact Q,K from num_kv_heads -> num_heads + - Delta rule state update: state [nheads, head_dim, head_dim] + - Gated RMSNorm with separate gate projection (in_proj_z) + - Decay: -exp(A_log) * softplus(in_proj_a(x) + dt_bias) per head + - Beta (write strength): sigmoid(in_proj_b(x)) per head + + Full attention layers: + - q_proj [2*attn_size, hidden] -> split query + gate + - QK-norm with (1+weight) centering + - Partial RoPE (partial_rotary_factor=0.25, 64 of 256 dims) + - KV cache + scaled dot-product attention + - Output gating: attn_out * sigmoid(gate) + +Weight key mapping (HF -> engine), verified against Qwen/Qwen3.8-27B: + model.language_model.embed_tokens.weight -> embedding + model.language_model.layers.{i}.input_layernorm.weight -> layer.{i}.input_norm + --- DeltaNet (linear_attention) layers --- + model.language_model.layers.{i}.linear_attn.in_proj_qkv.weight -> deltanet_in_proj_qkv + model.language_model.layers.{i}.linear_attn.in_proj_z.weight -> deltanet_z_proj (gate) + model.language_model.layers.{i}.linear_attn.in_proj_a.weight -> deltanet_a_proj (decay) + model.language_model.layers.{i}.linear_attn.in_proj_b.weight -> deltanet_b_proj (beta) + model.language_model.layers.{i}.linear_attn.A_log -> A + model.language_model.layers.{i}.linear_attn.dt_bias -> dt_bias + model.language_model.layers.{i}.linear_attn.conv1d.weight/bias -> conv1d + model.language_model.layers.{i}.linear_attn.norm.weight -> deltanet_norm + model.language_model.layers.{i}.linear_attn.out_proj.weight -> deltanet_out_proj + --- Full attention layers --- + model.language_model.layers.{i}.self_attn.q_proj.weight -> split: w_q + w_gate_attn + model.language_model.layers.{i}.self_attn.k_proj.weight -> w_k (keep compacted) + model.language_model.layers.{i}.self_attn.v_proj.weight -> w_v (keep compacted) + model.language_model.layers.{i}.self_attn.o_proj.weight -> w_o + model.language_model.layers.{i}.self_attn.q_norm.weight -> q_norm ((1+w) tiled) + model.language_model.layers.{i}.self_attn.k_norm.weight -> k_norm ((1+w) tiled) + --- SwiGLU MLP (both layer types) --- + model.language_model.layers.{i}.mlp.gate_proj.weight -> w_gate + model.language_model.layers.{i}.mlp.up_proj.weight -> w_up + model.language_model.layers.{i}.mlp.down_proj.weight -> w_down + model.language_model.layers.{i}.post_attention_layernorm.weight -> post_attn_norm + --- Final --- + model.language_model.norm.weight -> final_norm + lm_head.weight -> w_lm_head +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import numpy as np +from tensorrt_model_connect import trt_compat + +from .config import ModelConfig +from .checkpoint_mapper import ( + WeightDict, + _open_safetensors, + _load_tensor, + _has_tensor, + _transpose_2d, +) +from . import graph_ops +from . import graph_blocks + + +trt = trt_compat.get_trt() + + +def _parse_layer_types(raw_types: list[str]) -> list[str]: + """Normalize layer type strings to 'deltanet' or 'attention'.""" + mapping = { + "linear": "deltanet", + "linear_attention": "deltanet", + "full": "attention", + "full_attention": "attention", + } + return [mapping.get(t.lower(), t.lower()) for t in raw_types] + + +def _prepare_runtime_inputs( + network, + work_trt_dtype, + attention_mask, + conv_state_inputs, + ssm_state_inputs, + cache_k_inputs, + cache_v_inputs, +): + """Cast storage tensors while preserving DeltaNet recurrence in FP32.""" + if work_trt_dtype == trt.float32: + return ( + attention_mask, + conv_state_inputs, + ssm_state_inputs, + cache_k_inputs, + cache_v_inputs, + ) + + def cast_all(tensors): + return [ + network.add_cast(tensor, work_trt_dtype).get_output(0) + for tensor in tensors + ] + + return ( + network.add_cast(attention_mask, work_trt_dtype).get_output(0), + cast_all(conv_state_inputs), + # HF keeps the DeltaNet recurrent state in FP32. Never quantize this + # persistent input before the per-token recurrence. + ssm_state_inputs, + cast_all(cache_k_inputs), + cast_all(cache_v_inputs), + ) + + +class Qwen38Plugin: + name = "qwen3_8" + runtime_strategy = "qwen3_8_hybrid_mamba_attention" + + # Config keys that Qwen3.8 text configs carry and Qwen3.5 text configs do + # not. Qwen3.8 ships `model_type: "qwen3_5"` and + # `architectures: ["Qwen3_5ForConditionalGeneration"]`, so the checkpoint + # strings are useless for telling the two families apart and the config + # body is the only available signal. + _QWEN38_MARKER_KEYS = ("output_gate_type",) + # Conversely, every Qwen3.5 text config carries `mlp_only_layers`; Qwen3.8 + # dropped it. Requiring its absence keeps a future Qwen3.5 refresh that + # gained `output_gate_type` from being mis-claimed by this family. + _QWEN35_MARKER_KEYS = ("mlp_only_layers",) + + def matches(self, model_type: str) -> bool: + mt = model_type.lower() + return mt in {"qwen3_8", "qwen3.8", "qwen38"} + + def matches_config(self, config: object) -> bool: + """Claim Qwen3.8 checkpoints that masquerade as `model_type: qwen3_5`. + + Returning False here is what lets a genuine Qwen3.5 checkpoint fall + through to the qwen3_5 family: family discovery consults + `architecture_patterns` first, and only an affirmative + `matches_config` binds the model to this plugin. + """ + raw = getattr(config, "raw", None) + if not isinstance(raw, dict): + return False + + architectures = raw.get("architectures") or [] + if isinstance(architectures, str): + architectures = [architectures] + arch_ok = any( + "qwen3_5forconditionalgeneration" in str(value).lower() + for value in architectures + ) + model_type_ok = str(raw.get("model_type", "")).lower() in { + "qwen3_5", "qwen3.5", "qwen3_8", "qwen3.8", + } + if not arch_ok and not model_type_ok: + return False + + text_cfg = raw.get("text_config") + if not isinstance(text_cfg, dict): + text_cfg = raw + if any(key in text_cfg for key in self._QWEN35_MARKER_KEYS): + return False + return any(key in text_cfg for key in self._QWEN38_MARKER_KEYS) + + def load_weights( + self, model_dir: str, config: ModelConfig, + ) -> WeightDict: + model_dir_path = Path(model_dir) + readers = _open_safetensors(model_dir_path) + + hidden = config.hidden_size + vocab = config.vocab_size + num_layers = config.num_hidden_layers + raw = config.raw + + # Text config may be nested under text_config + text_cfg = raw.get("text_config", raw) + + # Parse layer types + raw_layer_types = text_cfg.get("layer_types", ["linear"] * num_layers) + layer_types = _parse_layer_types(raw_layer_types) + assert len(layer_types) == num_layers, ( + f"layer_types length {len(layer_types)} != num_hidden_layers {num_layers}") + + # Full attention dimensions + num_heads = config.num_attention_heads + num_kv_heads = config.num_key_value_heads + head_dim = config.head_dim + attn_size = num_heads * head_dim + kv_size = num_kv_heads * head_dim + + # DeltaNet dimensions (from text_config linear_* fields) + deltanet_num_heads = text_cfg.get("linear_num_value_heads", 32) + deltanet_num_kv_heads = text_cfg.get("linear_num_key_heads", 16) + deltanet_head_dim = text_cfg.get("linear_value_head_dim", + text_cfg.get("linear_key_head_dim", 128)) + d_inner = deltanet_num_heads * deltanet_head_dim + deltanet_qk_dim = deltanet_num_kv_heads * deltanet_head_dim + conv_dim = deltanet_qk_dim + deltanet_qk_dim + d_inner # Q + K + V + d_conv = text_cfg.get("linear_conv_kernel_dim", 4) + + # MLP dimensions + mlp_size = config.intermediate_size + + # RoPE config for full attention layers + # rope_parameters may be nested in text_config + rope_params = text_cfg.get("rope_parameters", {}) + partial_rotary_factor = rope_params.get( + "partial_rotary_factor", + text_cfg.get("partial_rotary_factor", 0.25)) + rope_theta = rope_params.get( + "rope_theta", + text_cfg.get("rope_theta", config.rope_theta)) + + weights = WeightDict() + + # Embedding + embed_key = "model.language_model.embed_tokens.weight" + if not _has_tensor(readers, embed_key): + embed_key = "model.embed_tokens.weight" + embedding = _load_tensor(readers, embed_key) + assert embedding.shape == (vocab, hidden), ( + f"Embedding shape {embedding.shape} != ({vocab}, {hidden})") + weights["embedding"] = embedding.astype(np.float32) + + deltanet_count = 0 + attn_count = 0 + + for layer_idx in range(num_layers): + lt = layer_types[layer_idx] + prefix = f"layer.{layer_idx}" + hf_prefix = f"model.language_model.layers.{layer_idx}" + + # Input layernorm (all layer types) + # Qwen3.8 uses (1+weight) centering in RMSNorm + norm_key = f"{hf_prefix}.input_layernorm.weight" + if _has_tensor(readers, norm_key): + weights[f"{prefix}.input_norm"] = ( + 1.0 + _load_tensor(readers, norm_key).astype(np.float32)) + else: + weights[f"{prefix}.input_norm"] = np.ones( + hidden, dtype=np.float32) + + # Post-attention layernorm (all layer types) + post_norm_key = f"{hf_prefix}.post_attention_layernorm.weight" + if _has_tensor(readers, post_norm_key): + weights[f"{prefix}.post_attn_norm"] = ( + 1.0 + _load_tensor(readers, post_norm_key).astype(np.float32)) + else: + weights[f"{prefix}.post_attn_norm"] = np.ones( + hidden, dtype=np.float32) + + if lt == "deltanet": + self._load_deltanet_weights( + readers, weights, prefix, hf_prefix, + hidden, d_inner, conv_dim, d_conv, + deltanet_num_heads, deltanet_num_kv_heads, + deltanet_head_dim) + deltanet_count += 1 + + elif lt == "attention": + self._load_attention_weights( + readers, weights, prefix, hf_prefix, + hidden, attn_size, kv_size, + num_heads, num_kv_heads, head_dim) + attn_count += 1 + + # SwiGLU MLP (all layer types) + self._load_mlp_weights( + readers, weights, prefix, hf_prefix, + hidden, mlp_size) + + # Final norm (also uses (1+weight) centering) + final_norm_key = "model.language_model.norm.weight" + if not _has_tensor(readers, final_norm_key): + final_norm_key = "model.norm.weight" + if _has_tensor(readers, final_norm_key): + weights["final_norm"] = ( + 1.0 + _load_tensor(readers, final_norm_key).astype(np.float32)) + else: + weights["final_norm"] = np.ones(hidden, dtype=np.float32) + + # LM head + lm_head_key = "lm_head.weight" + if _has_tensor(readers, lm_head_key): + weights["w_lm_head"] = _transpose_2d( + _load_tensor(readers, lm_head_key), "lm_head") + else: + weights["w_lm_head"] = _transpose_2d( + embedding.copy(), "embedding_tied") + + # Metadata for engine builder + weights["_layer_types"] = layer_types + weights["_d_inner"] = d_inner + weights["_d_conv"] = d_conv + weights["_conv_dim"] = conv_dim + weights["_deltanet_num_heads"] = deltanet_num_heads + weights["_deltanet_num_kv_heads"] = deltanet_num_kv_heads + weights["_deltanet_head_dim"] = deltanet_head_dim + weights["_num_mamba_layers"] = deltanet_count + weights["_num_attention_layers"] = attn_count + weights["_attn_size"] = attn_size + weights["_mlp_size"] = mlp_size + weights["_partial_rotary_factor"] = partial_rotary_factor + weights["_rope_theta"] = rope_theta + + return weights + + def _load_deltanet_weights( + self, readers, weights, prefix, hf_prefix, + hidden, d_inner, conv_dim, d_conv, + num_heads, num_kv_heads, head_dim, + ): + """Load DeltaNet (linear attention) layer weights.""" + attn_prefix = f"{hf_prefix}.linear_attn" + + # in_proj_qkv (QKV combined): [conv_dim, hidden] -> transpose + in_proj_raw = _load_tensor(readers, f"{attn_prefix}.in_proj_qkv.weight") + weights[f"{prefix}.deltanet_in_proj_qkv"] = _transpose_2d( + in_proj_raw, "deltanet_in_proj_qkv") + + # Gate projection (z): [d_inner, hidden] -> transpose + z_proj_raw = _load_tensor(readers, f"{attn_prefix}.in_proj_z.weight") + weights[f"{prefix}.deltanet_z_proj"] = _transpose_2d( + z_proj_raw, "deltanet_z_proj") + + # Decay projection (a): [num_heads, hidden] -> transpose + a_proj_raw = _load_tensor(readers, f"{attn_prefix}.in_proj_a.weight") + weights[f"{prefix}.deltanet_a_proj"] = _transpose_2d( + a_proj_raw, "deltanet_a_proj") + + # Beta projection (b): [num_heads, hidden] -> transpose + b_proj_raw = _load_tensor(readers, f"{attn_prefix}.in_proj_b.weight") + weights[f"{prefix}.deltanet_b_proj"] = _transpose_2d( + b_proj_raw, "deltanet_b_proj") + + # A_log: [num_heads] -> precompute -exp(A_log) + A_log = _load_tensor(readers, f"{attn_prefix}.A_log") + weights[f"{prefix}.A"] = -np.exp(A_log.astype(np.float32)) + + # dt_bias: [num_heads] + dt_bias = _load_tensor(readers, f"{attn_prefix}.dt_bias") + weights[f"{prefix}.dt_bias"] = dt_bias.astype(np.float32) + + # conv1d: [conv_dim, 1, d_conv] -> reshape to [conv_dim, d_conv] + conv_w = _load_tensor(readers, f"{attn_prefix}.conv1d.weight") + weights[f"{prefix}.conv1d_weight"] = conv_w.reshape( + conv_dim, d_conv).astype(np.float32) + + conv_b_key = f"{attn_prefix}.conv1d.bias" + if _has_tensor(readers, conv_b_key): + weights[f"{prefix}.conv1d_bias"] = _load_tensor( + readers, conv_b_key).astype(np.float32) + else: + weights[f"{prefix}.conv1d_bias"] = np.zeros( + conv_dim, dtype=np.float32) + + # Gated RMSNorm weight: [head_dim] -> tile to [d_inner] + norm_key = f"{attn_prefix}.norm.weight" + if _has_tensor(readers, norm_key): + norm_raw = _load_tensor(readers, norm_key).astype(np.float32) + # If weight is per-head (head_dim), tile to d_inner + if norm_raw.shape[0] == head_dim and head_dim < d_inner: + norm_raw = np.tile(norm_raw, num_heads) + weights[f"{prefix}.deltanet_norm"] = norm_raw + else: + weights[f"{prefix}.deltanet_norm"] = np.ones( + d_inner, dtype=np.float32) + + # Output projection: [hidden, d_inner] -> transpose + out_raw = _load_tensor(readers, f"{attn_prefix}.out_proj.weight") + weights[f"{prefix}.deltanet_out_proj"] = _transpose_2d( + out_raw, "deltanet_out_proj") + + def _load_attention_weights( + self, readers, weights, prefix, hf_prefix, + hidden, attn_size, kv_size, + num_heads, num_kv_heads, head_dim, + ): + """Load full self-attention layer weights.""" + attn_prefix = f"{hf_prefix}.self_attn" + + # q_proj: [2*attn_size, hidden] -> split per head into query + gate + # HF does: q_proj(x).view(B, seq, num_heads, 2*head_dim).chunk(2, dim=-1) + # This interleaves: for each head, first head_dim dims are query, next are gate + q_raw = _load_tensor(readers, f"{attn_prefix}.q_proj.weight") + # q_raw: [num_heads * 2 * head_dim, hidden] = [8192, 4096] + # Reshape to [num_heads, 2*head_dim, hidden], split, reshape back + q_reshaped = q_raw.reshape(num_heads, 2 * head_dim, hidden) + q_part = q_reshaped[:, :head_dim, :].reshape(attn_size, hidden) + gate_part = q_reshaped[:, head_dim:, :].reshape(attn_size, hidden) + weights[f"{prefix}.w_q"] = _transpose_2d(q_part, "q_proj") + weights[f"{prefix}.w_gate_attn"] = _transpose_2d(gate_part, "gate_proj") + + # k_proj: [kv_size, hidden] -> keep compact + k_raw = _load_tensor(readers, f"{attn_prefix}.k_proj.weight") + k_t = _transpose_2d(k_raw, "k_proj") + weights[f"{prefix}.w_k"] = k_t + + # v_proj: [kv_size, hidden] -> keep compact + v_raw = _load_tensor(readers, f"{attn_prefix}.v_proj.weight") + v_t = _transpose_2d(v_raw, "v_proj") + weights[f"{prefix}.w_v"] = v_t + + # o_proj: [hidden, attn_size] -> transpose + o_raw = _load_tensor(readers, f"{attn_prefix}.o_proj.weight") + weights[f"{prefix}.w_o"] = _transpose_2d(o_raw, "o_proj") + + # QK-norm with (1+weight) centering, tiled to num_heads + q_norm_key = f"{attn_prefix}.q_norm.weight" + if _has_tensor(readers, q_norm_key): + q_norm_raw = _load_tensor(readers, q_norm_key).astype(np.float32) + q_norm_centered = 1.0 + q_norm_raw # (1+weight) centering + weights[f"{prefix}.q_norm"] = np.tile( + q_norm_centered, num_heads) + k_norm_key = f"{attn_prefix}.k_norm.weight" + if _has_tensor(readers, k_norm_key): + k_norm_raw = _load_tensor(readers, k_norm_key).astype(np.float32) + k_norm_centered = 1.0 + k_norm_raw + weights[f"{prefix}.k_norm"] = np.tile( + k_norm_centered, num_kv_heads) + + def _load_mlp_weights( + self, readers, weights, prefix, hf_prefix, + hidden, mlp_size, + ): + """Load SwiGLU MLP weights.""" + gate_key = f"{hf_prefix}.mlp.gate_proj.weight" + up_key = f"{hf_prefix}.mlp.up_proj.weight" + down_key = f"{hf_prefix}.mlp.down_proj.weight" + + if _has_tensor(readers, gate_key): + weights[f"{prefix}.w_gate"] = _transpose_2d( + _load_tensor(readers, gate_key), "gate_proj") + weights[f"{prefix}.w_up"] = _transpose_2d( + _load_tensor(readers, up_key), "up_proj") + weights[f"{prefix}.w_down"] = _transpose_2d( + _load_tensor(readers, down_key), "down_proj") + + 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: + """Build hybrid TRT engine with DeltaNet + attention layers.""" + if quant_ctx is not None: + # This graph emits plain matmuls; it never threads a quantization + # context into its projections. Accepting quant_ctx silently would + # return an unquantized engine for a build the caller asked to + # quantize, so fail loudly instead. + raise NotImplementedError( + "Qwen3.8 does not support quantized builds; " + "build without --quantize/--fp8") + hidden = config.hidden_size + vocab = config.vocab_size + num_layers = config.num_hidden_layers + + layer_types: list[str] = weights["_layer_types"] + d_inner: int = weights["_d_inner"] + d_conv: int = weights["_d_conv"] + conv_dim: int = weights["_conv_dim"] + deltanet_num_heads: int = weights["_deltanet_num_heads"] + deltanet_num_kv_heads: int = weights["_deltanet_num_kv_heads"] + deltanet_head_dim: int = weights["_deltanet_head_dim"] + num_mamba: int = weights["_num_mamba_layers"] + num_attn: int = weights["_num_attention_layers"] + attn_size: int = weights["_attn_size"] + mlp_size: int = weights["_mlp_size"] + partial_rotary_factor: float = weights["_partial_rotary_factor"] + if precision == "fp16": + work_np_dtype, work_trt_dtype = np.float16, trt.float16 + elif precision == "fp32": + work_np_dtype, work_trt_dtype = np.float32, trt.float32 + else: + raise ValueError( + f"Unsupported Qwen3.8 precision {precision!r}; expected fp32 or fp16") + requested_fp32_layers = frozenset( + int(layer) for layer in config.raw.get("_fp32_layers", ())) + invalid_fp32_layers = sorted( + layer for layer in requested_fp32_layers + if layer < 0 or layer >= num_layers) + if invalid_fp32_layers: + raise ValueError( + "fp32_layers contains out-of-range indices: " + f"{invalid_fp32_layers}") + + num_heads = config.num_attention_heads + num_kv_heads = config.num_key_value_heads + head_dim = attn_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 + + 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() + + # --- 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, attention_window)) + + # DeltaNet state inputs (conv + ssm per DeltaNet layer) + conv_state_inputs = [] + ssm_state_inputs = [] + for mi in range(num_mamba): + cs = network.add_input( + graph_ops.layer_tensor_name("conv_state", mi), + trt.float32, (conv_dim, d_conv)) + ss = network.add_input( + graph_ops.layer_tensor_name("ssm_state", mi), + trt.float32, (deltanet_num_heads, deltanet_head_dim, deltanet_head_dim)) + conv_state_inputs.append(cs) + ssm_state_inputs.append(ss) + + # Attention KV cache inputs + cache_k_inputs = [] + cache_v_inputs = [] + for ai in range(num_attn): + ck = network.add_input( + graph_ops.layer_tensor_name("cache_k", ai), + work_trt_dtype, (max_cache_length, kv_attention_size)) + cv = network.add_input( + graph_ops.layer_tensor_name("cache_v", ai), + work_trt_dtype, (max_cache_length, kv_attention_size)) + cache_k_inputs.append(ck) + cache_v_inputs.append(cv) + + ( + attention_mask, + conv_state_inputs, + ssm_state_inputs, + cache_k_inputs, + cache_v_inputs, + ) = _prepare_runtime_inputs( + network, + work_trt_dtype, + attention_mask, + conv_state_inputs, + ssm_state_inputs, + cache_k_inputs, + cache_v_inputs, + ) + + # --- Shared constants --- + embedding_table = graph_ops.add_constant( + network, (vocab, hidden), weights["embedding"], + dtype=work_np_dtype) + eps_tensor = graph_ops.add_constant( + network, (1, 1), + np.array([config.rms_norm_eps], dtype=work_np_dtype), + dtype=work_np_dtype) + + rope_theta: float = weights["_rope_theta"] + rotary_embedding_dim = int(head_dim * partial_rotary_factor) + + # RoPE tables for full attention layers (partial rotary) + cos_half = graph_ops.make_rope_table_half_dim( + attention_window, head_dim, rope_theta, + cosine=True, partial_rotary_factor=partial_rotary_factor) + sin_half = graph_ops.make_rope_table_half_dim( + attention_window, head_dim, rope_theta, + cosine=False, partial_rotary_factor=partial_rotary_factor) + + cos_half_tensor = graph_ops.add_constant( + network, cos_half.shape, cos_half, dtype=work_np_dtype) + sin_half_tensor = graph_ops.add_constant( + network, sin_half.shape, sin_half, dtype=work_np_dtype) + + # --- Embedding --- + gather = network.add_gather(embedding_table, token_id, 0) + hidden_state = gather.get_output(0) + + if debug_layer_outputs: + _mark_debug_output(network, hidden_state, "debug_embed") + + # --- Layer stack --- + present_conv_outputs = [] + present_ssm_outputs = [] + present_k_outputs = [] + present_v_outputs = [] + mamba_counter = 0 + attn_counter = 0 + + for layer_idx in range(num_layers): + prefix = f"layer.{layer_idx}" + lt = layer_types[layer_idx] + layer_is_fp32 = ( + precision == "fp16" and layer_idx in requested_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 layer_cast(tensor): + if tensor.dtype == layer_trt_dtype: + return tensor + return network.add_cast( + tensor, layer_trt_dtype).get_output(0) + + if lt == "deltanet": + result = _add_deltanet_layer( + network=network, + hidden=layer_cast(hidden_state), + conv_state_in=layer_cast( + conv_state_inputs[mamba_counter]), + # Transformers casts the DeltaNet recurrence and its + # persistent state to FP32 even for FP16 checkpoints. + ssm_state_in=ssm_state_inputs[mamba_counter], + eps_tensor=layer_cast(eps_tensor), + weights=weights, + prefix=prefix, + hidden_size=hidden, + d_inner=d_inner, + d_conv=d_conv, + conv_dim=conv_dim, + num_heads=deltanet_num_heads, + num_kv_heads=deltanet_num_kv_heads, + head_dim=deltanet_head_dim, + mlp_size=mlp_size, + dtype=layer_np_dtype, + ) + hidden_state = result["hidden"] + present_conv_outputs.append(result["present_conv"]) + present_ssm_outputs.append(result["present_ssm"]) + mamba_counter += 1 + + elif lt == "attention": + result = _add_full_attention_layer( + network=network, + hidden=layer_cast(hidden_state), + cache_k=layer_cast(cache_k_inputs[attn_counter]), + cache_v=layer_cast(cache_v_inputs[attn_counter]), + attention_mask=layer_cast(attention_mask), + position_id=position_id, + cos_half_tensor=layer_cast(cos_half_tensor), + sin_half_tensor=layer_cast(sin_half_tensor), + eps_tensor=layer_cast(eps_tensor), + weights=weights, + prefix=prefix, + hidden_size=hidden, + attn_size=attn_size, + kv_attention_size=kv_attention_size, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + rotary_embedding_dim=rotary_embedding_dim, + max_cache_length=max_cache_length, + mlp_size=mlp_size, + dtype=layer_np_dtype, + ) + hidden_state = result["hidden"] + present_k_outputs.append(result["present_k"]) + present_v_outputs.append(result["present_v"]) + attn_counter += 1 + + if debug_layer_outputs: + _mark_debug_output( + network, hidden_state, f"debug_hidden_{layer_idx}") + + # --- Final norm --- + if hidden_state.dtype != work_trt_dtype: + hidden_state = network.add_cast( + hidden_state, work_trt_dtype).get_output(0) + final_norm = weights.get("final_norm") + if final_norm is not None and len(final_norm) > 0: + hidden_state = graph_ops.add_rms_norm( + network, hidden_state, hidden, final_norm, eps_tensor, + dtype=work_np_dtype) + + # --- LM head --- + logits = graph_ops.add_matmul_rhs_constant( + network, hidden_state, hidden, vocab, weights["w_lm_head"], + dtype=work_np_dtype) + b_out = np.zeros(vocab, dtype=work_np_dtype) + logits = graph_ops.add_bias_sum( + network, logits, vocab, b_out, dtype=work_np_dtype) + if logits.dtype != trt.float32: + logits = network.add_cast(logits, trt.float32).get_output(0) + logits.name = "logits" + network.mark_output(logits) + + # --- Present state outputs --- + for mi in range(num_mamba): + pc = present_conv_outputs[mi] + ps = present_ssm_outputs[mi] + if pc.dtype != trt.float32: + pc = network.add_cast(pc, trt.float32).get_output(0) + if ps.dtype != trt.float32: + ps = network.add_cast(ps, trt.float32).get_output(0) + pc.name = graph_ops.layer_tensor_name("present_conv", mi) + ps.name = graph_ops.layer_tensor_name("present_ssm", mi) + network.mark_output(pc) + network.mark_output(ps) + + for ai in range(num_attn): + pk = present_k_outputs[ai] + pv = present_v_outputs[ai] + if pk.dtype != work_trt_dtype: + pk = network.add_cast(pk, work_trt_dtype).get_output(0) + if pv.dtype != work_trt_dtype: + pv = network.add_cast(pv, work_trt_dtype).get_output(0) + pk.name = graph_ops.layer_tensor_name("present_k", ai) + pv.name = graph_ops.layer_tensor_name("present_v", ai) + network.mark_output(pk) + network.mark_output(pv) + + # --- Build --- + if verbose: + print(f"[trtmc build] Building Qwen3.8 hybrid TRT engine " + f"({num_layers} layers: {num_mamba} deltanet + " + f"{num_attn} attention, " + f"hidden={hidden}, d_inner={d_inner}, " + f"nheads_dn={deltanet_num_heads}, " + f"head_dim_dn={deltanet_head_dim}, " + f"cache={max_cache_length}) ...", + 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 get_bundle_config_overrides(self, config: ModelConfig) -> dict: + """Inject hybrid-specific config fields into the bundle.""" + raw = config.raw + text_cfg = raw.get("text_config", raw) + + raw_layer_types = text_cfg.get("layer_types", []) + layer_types = _parse_layer_types(raw_layer_types) + + deltanet_num_heads = text_cfg.get("linear_num_value_heads", 32) + deltanet_head_dim = text_cfg.get("linear_value_head_dim", + text_cfg.get("linear_key_head_dim", 128)) + deltanet_num_kv_heads = text_cfg.get("linear_num_key_heads", 16) + d_inner = deltanet_num_heads * deltanet_head_dim + d_conv = text_cfg.get("linear_conv_kernel_dim", 4) + deltanet_qk_dim = deltanet_num_kv_heads * deltanet_head_dim + conv_dim = deltanet_qk_dim + deltanet_qk_dim + d_inner + + num_mamba = sum(1 for lt in layer_types if lt == "deltanet") + num_attn = sum(1 for lt in layer_types if lt == "attention") + + # Qwen3.8 keeps every decoder dimension under `text_config`, but the + # C++ runtime reads the bundle config with a top-level nlohmann lookup + # (`extract_json_int` -> `j.find(key)`), not a recursive search. Left + # nested, `hidden_size`/`num_attention_heads`/`num_key_value_heads`/ + # `head_dim` all resolve to their fallbacks, `compute_kv_dim()` returns + # 0, and the KV cache allocates zero-sized tensors -- `ok()` is false + # and pipeline construction fails with "Failed to create Qwen38KvCache". + # Publishing flat copies here is the family-owned fix: bundle config + # overrides are emitted ahead of the raw config body, so the runtime + # sees real dimensions while `text_config` stays intact for the + # Python side. + # + # `eos_token_id` is deliberately not republished: the builder already + # emits a top-level list from generation_config.json, which carries all + # stop ids, whereas text_config holds only one. + flat_dims = {} + for key in ("vocab_size", "hidden_size", "num_hidden_layers", + "num_attention_heads", "num_key_value_heads", "head_dim", + "intermediate_size", "max_position_embeddings", + "rms_norm_eps", "bos_token_id"): + value = text_cfg.get(key) + if value is not None: + flat_dims[key] = value + if "head_dim" not in flat_dims: + heads = flat_dims.get("num_attention_heads", 0) + hidden = flat_dims.get("hidden_size", 0) + if heads and hidden: + flat_dims["head_dim"] = hidden // heads + + return { + **flat_dims, + "layer_types": layer_types, + "num_mamba_layers": num_mamba, + "num_attention_layers": num_attn, + "d_inner": d_inner, + "mamba_d_state": deltanet_head_dim, + "mamba_d_conv": d_conv, + "mamba_nheads": deltanet_num_heads, + "mamba_head_dim": deltanet_head_dim, + "conv_dim": conv_dim, + } + + +def _mark_debug_output( + network: trt.INetworkDefinition, + tensor: trt.ITensor, + name: str, +) -> None: + identity = network.add_identity(tensor) + cast = network.add_cast(identity.get_output(0), trt.float32) + out = cast.get_output(0) + out.name = name + network.mark_output(out) + + +def _add_deltanet_layer( + *, + network: trt.INetworkDefinition, + hidden: trt.ITensor, + conv_state_in: trt.ITensor, + ssm_state_in: trt.ITensor, + eps_tensor: trt.ITensor, + weights: WeightDict, + prefix: str, + hidden_size: int, + d_inner: int, + d_conv: int, + conv_dim: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + mlp_size: int, + dtype: np.dtype = np.float32, +) -> dict[str, trt.ITensor]: + """Add one Gated DeltaNet layer (single-step decode). + + DeltaNet uses delta-rule linear attention with: + - Conv1d on QKV projection output + - L2-normalized Q,K with Compact GQA/MQA K/V + - Per-head decay (A_log + softplus(a + dt_bias)) + - Per-head beta (write strength via sigmoid) + - Delta rule state update: S' = decay*S + outer(k, (v - S@k)*beta) + - Gated output: norm(S'@q) * silu(z) * norm_weight + + Returns: {hidden, present_conv, present_ssm} + """ + qk_dim = num_kv_heads * head_dim # Q and K dimension before Compact GQA/MQA K/V + + # ===== 1. RMSNorm ===== + normed = graph_ops.add_rms_norm( + network, hidden, hidden_size, + weights[f"{prefix}.input_norm"], eps_tensor, dtype=dtype) + + # ===== 2. Input projections ===== + # QKV combined: [1, hidden] -> [1, conv_dim] + qkv = graph_ops.add_matmul_rhs_constant( + network, normed, hidden_size, conv_dim, + weights[f"{prefix}.deltanet_in_proj_qkv"], dtype=dtype) + + # Gate (z): [1, hidden] -> [1, d_inner] + z = graph_ops.add_matmul_rhs_constant( + network, normed, hidden_size, d_inner, + weights[f"{prefix}.deltanet_z_proj"], dtype=dtype) + + # Decay projection (a): [1, hidden] -> [1, num_heads] + a_raw = graph_ops.add_matmul_rhs_constant( + network, normed, hidden_size, num_heads, + weights[f"{prefix}.deltanet_a_proj"], dtype=dtype) + + # Beta projection (b): [1, hidden] -> [1, num_heads] + b_raw = graph_ops.add_matmul_rhs_constant( + network, normed, hidden_size, num_heads, + weights[f"{prefix}.deltanet_b_proj"], dtype=dtype) + + # ===== 3. Conv1d step on QKV ===== + # conv_state_in: [conv_dim, d_conv] + # qkv: [1, conv_dim] -> [conv_dim, 1] + qkv_col = network.add_shuffle(qkv) + qkv_col.reshape_dims = (conv_dim, 1) + + if d_conv > 1: + slice_layer = network.add_slice( + conv_state_in, + start=(0, 1), + shape=(conv_dim, d_conv - 1), + stride=(1, 1)) + new_conv_state = network.add_concatenation( + [slice_layer.get_output(0), qkv_col.get_output(0)]) + new_conv_state.axis = 1 + present_conv = new_conv_state.get_output(0) + else: + present_conv = qkv_col.get_output(0) + + conv_w = graph_ops.add_constant( + network, (conv_dim, d_conv), weights[f"{prefix}.conv1d_weight"], + dtype=dtype) + conv_prod = network.add_elementwise( + present_conv, conv_w, trt.ElementWiseOperation.PROD) + conv_sum = network.add_reduce( + conv_prod.get_output(0), trt.ReduceOperation.SUM, + 1 << 1, keep_dims=True) + conv_flat = network.add_shuffle(conv_sum.get_output(0)) + conv_flat.reshape_dims = (1, conv_dim) + conv_out = graph_ops.add_bias_sum( + network, conv_flat.get_output(0), conv_dim, + weights[f"{prefix}.conv1d_bias"], dtype=dtype) + qkv_activated = graph_ops.add_activation( + network, conv_out, "silu", dtype=dtype) + + # ===== 4. Split Q, K, V from activated output ===== + offset = 0 + q_slice = network.add_slice( + qkv_activated, start=(0, offset), shape=(1, qk_dim), stride=(1, 1)) + q_raw_t = q_slice.get_output(0) + offset += qk_dim + + k_slice = network.add_slice( + qkv_activated, start=(0, offset), shape=(1, qk_dim), stride=(1, 1)) + k_raw_t = k_slice.get_output(0) + offset += qk_dim + + v_slice = network.add_slice( + qkv_activated, start=(0, offset), shape=(1, d_inner), stride=(1, 1)) + v_raw = v_slice.get_output(0) + + # ===== 5. L2-normalize Q and K ===== + # Reshape to [num_kv_heads, head_dim], normalize per-head, reshape back + q_heads_in = network.add_shuffle(q_raw_t) + q_heads_in.reshape_dims = (num_kv_heads, head_dim) + q_normed = graph_ops.add_l2_norm( + network, q_heads_in.get_output(0), 1, eps=1e-6, dtype=dtype) + + k_heads_in = network.add_shuffle(k_raw_t) + k_heads_in.reshape_dims = (num_kv_heads, head_dim) + k_normed = graph_ops.add_l2_norm( + network, k_heads_in.get_output(0), 1, eps=1e-6, dtype=dtype) + + # ===== 6. keep compact Q,K from num_kv_heads -> num_heads ===== + heads_per_group = num_heads // num_kv_heads + + if heads_per_group > 1: + # Q: [num_kv_heads, head_dim] -> [num_kv_heads, 1, head_dim] -> + # tile -> [num_kv_heads, heads_per_group, head_dim] -> + # [num_heads, head_dim] + q_3d = network.add_shuffle(q_normed) + q_3d.reshape_dims = (num_kv_heads, 1, head_dim) + tile_ones = graph_ops.add_constant( + network, (1, heads_per_group, 1), + np.ones((1, heads_per_group, 1), dtype=dtype), dtype=dtype) + q_tiled = network.add_elementwise( + q_3d.get_output(0), tile_ones, trt.ElementWiseOperation.PROD) + q_expanded_s = network.add_shuffle(q_tiled.get_output(0)) + q_expanded_s.reshape_dims = (num_heads, head_dim) + q_expanded = q_expanded_s.get_output(0) + + k_3d = network.add_shuffle(k_normed) + k_3d.reshape_dims = (num_kv_heads, 1, head_dim) + k_tiled = network.add_elementwise( + k_3d.get_output(0), tile_ones, trt.ElementWiseOperation.PROD) + k_t_s = network.add_shuffle(k_tiled.get_output(0)) + k_t_s.reshape_dims = (num_heads, head_dim) + k_t = k_t_s.get_output(0) + else: + q_expanded = q_normed + k_t = k_normed + + # V: [1, d_inner] -> [num_heads, head_dim] + v_heads = network.add_shuffle(v_raw) + v_heads.reshape_dims = (num_heads, head_dim) + v_t = v_heads.get_output(0) + + # ===== 7. Compute decay: -exp(A_log) * softplus(a + dt_bias) per head ===== + # Transformers performs the decay and recurrent rule in FP32 even when + # the model projections use FP16. Keeping these tensors in the model + # storage dtype quantizes the persistent state again on every token. + recurrent_dtype = trt.float32 + + def recurrent_cast(tensor: trt.ITensor) -> trt.ITensor: + if tensor.dtype == recurrent_dtype: + return tensor + return network.add_cast(tensor, recurrent_dtype).get_output(0) + + # A: [num_heads] (precomputed as -exp(A_log)) + A_const = graph_ops.add_constant( + network, (1, num_heads), weights[f"{prefix}.A"], dtype=np.float32) + + # dt_bias: [num_heads] + dt_bias_const = graph_ops.add_constant( + network, (1, num_heads), weights[f"{prefix}.dt_bias"], dtype=np.float32) + a_biased = network.add_elementwise( + recurrent_cast(a_raw), dt_bias_const, trt.ElementWiseOperation.SUM) + + # softplus(a + dt_bias): log(1 + exp(x)) + a_exp = network.add_unary(a_biased.get_output(0), trt.UnaryOperation.EXP) + one = graph_ops.add_constant( + network, (1, 1), np.array([1.0], dtype=np.float32), dtype=np.float32) + a_exp_p1 = network.add_elementwise( + a_exp.get_output(0), one, trt.ElementWiseOperation.SUM) + a_softplus = network.add_unary( + a_exp_p1.get_output(0), trt.UnaryOperation.LOG) + + # decay = A * softplus(...) per head: [1, num_heads] + decay_flat = network.add_elementwise( + A_const, a_softplus.get_output(0), trt.ElementWiseOperation.PROD) + # exp(decay) for the state update: [1, num_heads] -> [num_heads, 1, 1] + decay_reshaped = network.add_shuffle(decay_flat.get_output(0)) + decay_reshaped.reshape_dims = (num_heads, 1, 1) + decay_exp = network.add_unary( + decay_reshaped.get_output(0), trt.UnaryOperation.EXP) + + # ===== 8. Compute beta: sigmoid(b) per head ===== + # b_raw: [1, num_heads] + beta = network.add_activation(b_raw, trt.ActivationType.SIGMOID) + # [1, num_heads] -> [num_heads, 1] + beta_reshaped = network.add_shuffle(recurrent_cast(beta.get_output(0))) + beta_reshaped.reshape_dims = (num_heads, 1) + + # ===== 9. Delta rule state update ===== + # HF state layout: [H, K_dim, V_dim] + # ssm_state_in: [num_heads, head_dim, head_dim] (K on axis -2, V on axis -1) + # k: [num_heads, head_dim], q: [num_heads, head_dim], v: [num_heads, head_dim] + + # 9a. Decay state first: state = state * exp(g) + decayed_state = network.add_elementwise( + decay_exp.get_output(0), recurrent_cast(ssm_state_in), + trt.ElementWiseOperation.PROD) + + # 9b. kv_mem = state^T @ k: read old value for this key + # [H, V, K] @ [H, K, 1] = [H, V, 1] (transpose state to swap K/V axes) + k_recurrent = recurrent_cast(k_t) + v_recurrent = recurrent_cast(v_t) + q_recurrent = recurrent_cast(q_expanded) + + k_col = network.add_shuffle(k_recurrent) + k_col.reshape_dims = (num_heads, head_dim, 1) + kv_old_3d = network.add_matrix_multiply( + decayed_state.get_output(0), trt.MatrixOperation.TRANSPOSE, + k_col.get_output(0), trt.MatrixOperation.NONE) + kv_old = network.add_shuffle(kv_old_3d.get_output(0)) + kv_old.reshape_dims = (num_heads, head_dim) + + # 9c. delta = (v - kv_mem) * beta + v_minus_old = network.add_elementwise( + v_recurrent, kv_old.get_output(0), trt.ElementWiseOperation.SUB) + v_delta = network.add_elementwise( + v_minus_old.get_output(0), beta_reshaped.get_output(0), + trt.ElementWiseOperation.PROD) + + # 9d. state_new = decayed_state + outer(k, delta) + # outer: k[:, :, None] * delta[:, None, :] = [H, K, 1] @ [H, 1, V] = [H, K, V] + k_col2 = network.add_shuffle(k_recurrent) + k_col2.reshape_dims = (num_heads, head_dim, 1) + v_delta_row = network.add_shuffle(v_delta.get_output(0)) + v_delta_row.reshape_dims = (num_heads, 1, head_dim) + outer_prod = network.add_matrix_multiply( + k_col2.get_output(0), trt.MatrixOperation.NONE, + v_delta_row.get_output(0), trt.MatrixOperation.NONE) + + new_state = network.add_elementwise( + decayed_state.get_output(0), outer_prod.get_output(0), + trt.ElementWiseOperation.SUM) + present_ssm = new_state.get_output(0) + + # 9e. output = state_new^T @ (q * scale) + # HF applies: query *= 1/sqrt(k_dim) + q_scale = graph_ops.add_constant( + network, (1, 1), + np.array([1.0 / np.sqrt(head_dim)], dtype=np.float32), dtype=np.float32) + q_scaled = network.add_elementwise( + q_recurrent, q_scale, trt.ElementWiseOperation.PROD) + # [H, V, K] @ [H, K, 1] = [H, V, 1] + q_col = network.add_shuffle(q_scaled.get_output(0)) + q_col.reshape_dims = (num_heads, head_dim, 1) + output_3d = network.add_matrix_multiply( + present_ssm, trt.MatrixOperation.TRANSPOSE, + q_col.get_output(0), trt.MatrixOperation.NONE) + output_flat = network.add_shuffle(output_3d.get_output(0)) + output_flat.reshape_dims = (1, d_inner) + + # ===== 10. Gated RMSNorm per-head: weight * norm(output) * silu(z) ===== + # The reference recurrent kernel returns the attention output in the model + # storage dtype before Qwen3_5RMSNormGated casts it back to FP32. + recurrent_output = output_flat.get_output(0) + if recurrent_output.dtype != hidden.dtype: + recurrent_output = network.add_cast( + recurrent_output, hidden.dtype).get_output(0) + + # HF norm operates per head_v_dim: reshape to [num_heads, head_dim], norm, reshape back + deltanet_norm_w = weights[f"{prefix}.deltanet_norm"] + # Use same eps as HF Qwen3_5RMSNormGated (config.rms_norm_eps = 1e-6) + eps_small = graph_ops.add_constant( + network, (1, 1), + np.array([1e-6], dtype=np.float32), dtype=np.float32) + + # Reshape output and z to [num_heads, head_dim] for per-head norm + output_heads = network.add_shuffle(recurrent_output) + output_heads.reshape_dims = (num_heads, head_dim) + norm_input = output_heads.get_output(0) + norm_output_dtype = norm_input.dtype + if dtype != np.float32: + norm_input = network.add_cast(norm_input, trt.float32).get_output(0) + + # Per-head RMSNorm: norm each head independently + sq = network.add_elementwise( + norm_input, norm_input, + 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_small, 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( + norm_input, recip.get_output(0), + trt.ElementWiseOperation.PROD) + + # Reshape back and apply weight + norm_flat = network.add_shuffle(normalized.get_output(0)) + norm_flat.reshape_dims = (1, d_inner) + gamma_t = graph_ops.add_constant( + network, (1, d_inner), deltanet_norm_w, dtype=np.float32) + normed_output = network.add_elementwise( + norm_flat.get_output(0), gamma_t, trt.ElementWiseOperation.PROD) + normed_output_tensor = normed_output.get_output(0) + if normed_output_tensor.dtype != norm_output_dtype: + normed_output_tensor = network.add_cast( + normed_output_tensor, norm_output_dtype).get_output(0) + + # Gate: multiply by silu(z) + z_activated = graph_ops.add_activation(network, z, "silu", dtype=dtype) + gated = network.add_elementwise( + normed_output_tensor, z_activated, + trt.ElementWiseOperation.PROD) + + # ===== 11. Output projection + residual ===== + out = graph_ops.add_matmul_rhs_constant( + network, gated.get_output(0), d_inner, hidden_size, + weights[f"{prefix}.deltanet_out_proj"], dtype=dtype) + + residual = network.add_elementwise( + hidden, out, trt.ElementWiseOperation.SUM) + hidden_after_attn = residual.get_output(0) + + # ===== 12. Post-attention norm + SwiGLU MLP + residual ===== + post_normed = graph_ops.add_rms_norm( + network, hidden_after_attn, hidden_size, + weights[f"{prefix}.post_attn_norm"], eps_tensor, dtype=dtype) + + mlp_out = graph_blocks.add_swiglu_mlp( + network, post_normed, + weights=weights, + prefix=prefix, + hidden_size=hidden_size, + mlp_size=mlp_size, + dtype=dtype, + ) + + mlp_residual = network.add_elementwise( + hidden_after_attn, mlp_out, + trt.ElementWiseOperation.SUM) + + return { + "hidden": mlp_residual.get_output(0), + "present_conv": present_conv, + "present_ssm": present_ssm, + } + + +def _add_full_attention_layer( + *, + network: trt.INetworkDefinition, + hidden: trt.ITensor, + cache_k: trt.ITensor, + cache_v: trt.ITensor, + attention_mask: trt.ITensor, + position_id: trt.ITensor, + cos_half_tensor: trt.ITensor, + sin_half_tensor: trt.ITensor, + eps_tensor: trt.ITensor, + weights: WeightDict, + prefix: str, + hidden_size: int, + attn_size: int, + kv_attention_size: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + rotary_embedding_dim: int, + max_cache_length: int, + mlp_size: int, + dtype: np.dtype = np.float32, +) -> dict[str, trt.ITensor]: + """Add one full self-attention layer with output gating. + + Qwen3.8 full attention has: + - QK-norm with (1+weight) centering + - Partial RoPE (25% of dims) + - Output gating: context * sigmoid(gate) BEFORE o_proj + - SwiGLU MLP after attention + + Returns: {hidden, present_k, present_v} + """ + attention_window = max_cache_length + 1 + + # Pre-attention norm + normed = graph_blocks.apply_norm( + network, hidden, hidden_size, + weights[f"{prefix}.input_norm"], + weights.get(f"{prefix}.input_norm_beta"), + eps_tensor, "rmsnorm", dtype=dtype) + + # QKV projections + q = graph_ops.add_matmul_rhs_constant( + network, normed, hidden_size, attn_size, + weights[f"{prefix}.w_q"], dtype=dtype) + k = graph_ops.add_matmul_rhs_constant( + network, normed, hidden_size, kv_attention_size, + weights[f"{prefix}.w_k"], dtype=dtype) + v = graph_ops.add_matmul_rhs_constant( + network, normed, hidden_size, kv_attention_size, + weights[f"{prefix}.w_v"], dtype=dtype) + + # Per-head QK 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) + + # Native RoPE + q = graph_ops.add_apply_rope_native( + network, q, num_heads, head_dim, cos_half_tensor, sin_half_tensor, + position_id, rotary_embedding_dim) + k = graph_ops.add_apply_rope_native( + network, k, num_kv_heads, head_dim, cos_half_tensor, sin_half_tensor, + position_id, rotary_embedding_dim) + + # Save present K/V + present_k = k + present_v = v + + # Reshape 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 + + mask_4d = graph_ops.add_2d_mask_to_4d(network, attention_mask) + context_flat = 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=attention_window, + mask=mask_4d) + + # Gate: applied BEFORE o_proj (HF order) + gate_attn_w = weights.get(f"{prefix}.w_gate_attn") + attn_out = context_flat + if gate_attn_w is not None: + gate = graph_ops.add_matmul_rhs_constant( + network, normed, hidden_size, attn_size, gate_attn_w, + dtype=dtype) + gate_sigmoid = network.add_activation(gate, trt.ActivationType.SIGMOID) + gated = network.add_elementwise( + attn_out, gate_sigmoid.get_output(0), + trt.ElementWiseOperation.PROD) + attn_out = gated.get_output(0) + + # Output projection (AFTER gate) + attn_out = graph_ops.add_matmul_rhs_constant( + network, attn_out, attn_size, hidden_size, + weights[f"{prefix}.w_o"], dtype=dtype) + + # Residual after attention + residual = network.add_elementwise( + hidden, attn_out, trt.ElementWiseOperation.SUM) + hidden_after_attn = residual.get_output(0) + + # Post-attention norm + SwiGLU MLP + residual + post_normed = graph_ops.add_rms_norm( + network, hidden_after_attn, hidden_size, + weights[f"{prefix}.post_attn_norm"], eps_tensor, dtype=dtype) + + mlp_out = graph_blocks.add_swiglu_mlp( + network, post_normed, + weights=weights, + prefix=prefix, + hidden_size=hidden_size, + mlp_size=mlp_size, + dtype=dtype, + ) + + mlp_residual = network.add_elementwise( + hidden_after_attn, mlp_out, + trt.ElementWiseOperation.SUM) + + return { + "hidden": mlp_residual.get_output(0), + "present_k": present_k, + "present_v": present_v, + } + + +plugin = Qwen38Plugin() diff --git a/src/runtime/domains/recurrent/README.md b/src/runtime/domains/recurrent/README.md index 4644c797d..bb45016a3 100644 --- a/src/runtime/domains/recurrent/README.md +++ b/src/runtime/domains/recurrent/README.md @@ -9,5 +9,6 @@ Current recurrent contract owners: - `src/runtime/models/rwkv/rwkv_recurrent_step_contracts.h` - `src/runtime/models/nemotron_h/nemotron_h_recurrent_step_contracts.h` - `src/runtime/models/qwen3_5/qwen3_5_recurrent_step_contracts.h` +- `src/runtime/models/qwen3_8/qwen3_8_recurrent_step_contracts.h` diff --git a/src/runtime/models/qwen3_8/MODEL.toml b/src/runtime/models/qwen3_8/MODEL.toml new file mode 100644 index 000000000..11c4c3e6d --- /dev/null +++ b/src/runtime/models/qwen3_8/MODEL.toml @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +id = "qwen3_8" +runtime_library = "libtrtmc_model_qwen3_8.so" +runtime_plugins = ["plugin.cpp|register_qwen3_8_plugin"] +runtime_strategies = ["qwen3_8_hybrid_mamba_attention"] +runtime_tests = [ + "test_qwen3_8_runtime_config_contract|test_qwen3_8_runtime_config_contract.cpp|_|_|_", + "test_qwen3_8_recurrent_output_initializers|test_qwen3_8_recurrent_output_initializers.cpp|_|_|_", + "test_qwen3_8_recurrent_pipeline|test_qwen3_8_recurrent_pipeline.cpp|trtmc_model_qwen3_8,trtmc_backend_trt|_|REQUIRES_TRT,REQUIRES_GPU", +] + +[validation_profiles] +decoder_debug = ["qwen3_8_hybrid_mamba_attention"] diff --git a/src/runtime/models/qwen3_8/chat_templates.cpp b/src/runtime/models/qwen3_8/chat_templates.cpp new file mode 100644 index 000000000..e8ffddd2e --- /dev/null +++ b/src/runtime/models/qwen3_8/chat_templates.cpp @@ -0,0 +1,50 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/models/qwen3_8/chat_templates.h" + +#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_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; +} + +} // namespace + +std::string qwen3_8_detect_chat_template_format(const std::string& jinja_template) { + if (jinja_template.empty()) + return {}; + if (jinja_template.find("<|im_start|>") != std::string::npos) + return "chatml"; + if (jinja_template.find("") != std::string::npos) + return "nemotron_h"; + return {}; +} + +std::string qwen3_8_apply_chat_template(const std::string& format, const std::string& prompt, + bool enable_thinking) { + if (format.empty()) + return prompt; + if (format == "chatml") + return apply_chatml(prompt, enable_thinking); + if (format == "nemotron_h") + return apply_nemotron_h(prompt, enable_thinking); + return prompt; +} + +} // namespace trtmc diff --git a/src/runtime/models/qwen3_8/chat_templates.h b/src/runtime/models/qwen3_8/chat_templates.h new file mode 100644 index 000000000..cc01a2e25 --- /dev/null +++ b/src/runtime/models/qwen3_8/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 qwen3_8_detect_chat_template_format(const std::string& jinja_template); +std::string qwen3_8_apply_chat_template(const std::string& format, const std::string& prompt, + bool enable_thinking = true); + +} // namespace trtmc diff --git a/src/runtime/models/qwen3_8/hybrid_state.cpp b/src/runtime/models/qwen3_8/hybrid_state.cpp new file mode 100644 index 000000000..fd2703b89 --- /dev/null +++ b/src/runtime/models/qwen3_8/hybrid_state.cpp @@ -0,0 +1,62 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/models/qwen3_8/hybrid_state.h" + +#include +#include + +namespace trtmc { + +// Every method below delegates to both members unconditionally, so null is +// rejected here rather than tolerated. That keeps ok() a health check on the +// device allocations instead of a null test callers may or may not have run. +Qwen38HybridState::Qwen38HybridState(std::unique_ptr kv, + std::unique_ptr ssm) + : kv_(std::move(kv)), ssm_(std::move(ssm)) { + if (!kv_ || !ssm_) + throw std::invalid_argument("Qwen38HybridState requires non-null KV and recurrent state"); +} + +void Qwen38HybridState::reset() { + kv_->reset(); + ssm_->reset(); +} + +void Qwen38HybridState::bind_to(TrtModule& module) { + kv_->bind_to(module); + ssm_->bind_to(module); +} + +void Qwen38HybridState::prepare_step(TensorMap& inputs, int32_t seq_len) { + kv_->prepare_step(inputs, seq_len); +} + +void Qwen38HybridState::advance(int32_t n_tokens) { + kv_->advance(n_tokens); + ssm_->advance(n_tokens); +} + +int32_t Qwen38HybridState::position() const { + return kv_->position(); +} + +int32_t Qwen38HybridState::max_length() const { + return kv_->max_length(); +} + +int32_t Qwen38HybridState::num_layers() const { + return kv_->num_layers() + ssm_->num_layers(); +} + +std::size_t Qwen38HybridState::device_memory_bytes() const { + return kv_->device_memory_bytes() + ssm_->device_memory_bytes(); +} + +bool Qwen38HybridState::ok() const { + return kv_->ok() && ssm_->ok(); +} + +} // namespace trtmc diff --git a/src/runtime/models/qwen3_8/hybrid_state.h b/src/runtime/models/qwen3_8/hybrid_state.h new file mode 100644 index 000000000..592bda6a0 --- /dev/null +++ b/src/runtime/models/qwen3_8/hybrid_state.h @@ -0,0 +1,42 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "runtime/models/qwen3_8/inference_state.h" +#include "runtime/models/qwen3_8/kv_cache.h" +#include "runtime/models/qwen3_8/recurrent_state.h" + +#include + +namespace trtmc { + +class Qwen38HybridState final : public Qwen38InferenceState { + public: + Qwen38HybridState(std::unique_ptr kv, std::unique_ptr ssm); + + 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; + int32_t max_length() const override; + int32_t num_layers() const override; + bool needs_attention_mask() const override { return true; } + std::size_t device_memory_bytes() const override; + const char* state_type() const override { return "qwen3_8_hybrid_kv_recurrent"; } + bool ok() const override; + + Qwen38KvCache* kv_cache() { return kv_.get(); } + const Qwen38KvCache* kv_cache() const { return kv_.get(); } + Qwen38RecurrentState* recurrent_state() { return ssm_.get(); } + const Qwen38RecurrentState* recurrent_state() const { return ssm_.get(); } + + private: + std::unique_ptr kv_; + std::unique_ptr ssm_; +}; + +} // namespace trtmc diff --git a/src/runtime/models/qwen3_8/inference_state.h b/src/runtime/models/qwen3_8/inference_state.h new file mode 100644 index 000000000..77671aa8d --- /dev/null +++ b/src/runtime/models/qwen3_8/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 + +// Qwen38InferenceState: 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: +// Qwen38KvCache — dense append-only (current default) +// Family-owned recurrent state — recurrent tensor state +// Family-owned hybrid state — Qwen38KvCache + 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 Qwen38InferenceState { + public: + virtual ~Qwen38InferenceState() = 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. + // Qwen38KvCache -> 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/qwen3_8/kv_cache.cpp b/src/runtime/models/qwen3_8/kv_cache.cpp new file mode 100644 index 000000000..dde76652c --- /dev/null +++ b/src/runtime/models/qwen3_8/kv_cache.cpp @@ -0,0 +1,399 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/models/qwen3_8/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); +} + +} // namespace + +Qwen38KvCache::Qwen38KvCache(int32_t num_layers, int32_t max_length, int32_t kv_dim, + cudaStream_t stream, DType cache_dtype, Qwen38KvCacheNames 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); + } + } + + 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); + } + + // Pre-allocate mask buffer: [max_length + 1] for dense causal mask. + mask_buf_.resize(static_cast(max_length) + 1); + + reset(); +} + +// Masked score constant is model-local. +static constexpr float kMaskedScore = -1.0e4F; + +void Qwen38KvCache::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 Qwen38KvCache::preferred_cache_rows() const { + if (!dynamic_binding_enabled_) + return max_length_; + return round_up_rows(std::max(position_, 1), kRuntimeBucketRows, max_length_); +} + +void Qwen38KvCache::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 Qwen38KvCache::mask_shape_for_engine(int32_t mask_width) 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}; + // Logical width, not mask_buf_.size(): the buffer is grown by a batched + // prefill and never shrunk, so its size overstates a later decode mask. + return {static_cast(mask_width)}; +} + +void Qwen38KvCache::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 Qwen38KvCache::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 Qwen38KvCache::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 Qwen38KvCache::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_t.dtype = DType::kFloat32; + inputs[names_.attention_mask] = mask_t; +} + +void Qwen38KvCache::prepare_step(TensorMap& inputs, int32_t seq_len) { + if (seq_len <= 0) + seq_len = 1; + write_position_input(inputs, seq_len); + if (seq_len > 1) + write_batched_mask(inputs, seq_len); + else + write_decode_mask(inputs); +} + +void Qwen38KvCache::prepare_bidirectional_step(TensorMap& inputs, int32_t seq_len) { + if (seq_len <= 0) + seq_len = 1; + write_position_input(inputs, seq_len); + write_bidirectional_mask(inputs, seq_len); +} + +void Qwen38KvCache::bind_to(TrtModule& module) { + bound_module_ = &module; + has_position_input_ = module.has_input(names_.position_id); + // 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 Qwen38KvCache::bind_cache_inputs(TrtModule& module) { + has_position_input_ = module.has_input(names_.position_id); + for (int32_t i = 0; i < num_layers_; ++i) { + 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()); + } +} + +void Qwen38KvCache::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("Qwen38KvCache::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( + "Qwen38KvCache::write_prefill_kv: per-layer pointer count mismatch"); + } + 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 Qwen38KvCache::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("Qwen38KvCache::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( + "Qwen38KvCache::append_prefill_kv: per-layer pointer count mismatch"); + } + 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 Qwen38KvCache::set_position(int32_t position) { + position_ = std::max(0, std::min(position, max_length_)); +} + +void Qwen38KvCache::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 && "Qwen38KvCache::advance: only n_tokens==1 supported"); + (void)n_tokens; + + // 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; + // src and dst overlap inside the same allocation, and cudaMemcpyAsync + // is undefined for overlapping ranges, so stage through scratch. One + // buffer serves every layer because all copies are ordered on stream_. + DeviceTensor& scratch_tensor = shift_scratch(); + if (shift_bytes > 0 && !scratch_tensor.ok()) + throw std::runtime_error( + "Qwen38KvCache: failed to allocate scratch for the cache shift"); + auto* scratch = static_cast(scratch_tensor.data()); + 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()); + if (shift_bytes > 0) { + cudaMemcpyAsync(scratch, ck + row_bytes, shift_bytes, cudaMemcpyDeviceToDevice, + stream_); + cudaMemcpyAsync(ck, scratch, shift_bytes, cudaMemcpyDeviceToDevice, stream_); + cudaMemcpyAsync(scratch, cv + row_bytes, shift_bytes, cudaMemcpyDeviceToDevice, + stream_); + cudaMemcpyAsync(cv, scratch, 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 Qwen38KvCache::reset() { + // Reset only the logical sequence length. Attention masks hide every + // stale cache row, and each present row is overwritten before use. + position_ = 0; +} + +std::size_t Qwen38KvCache::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; +} + +DeviceTensor& Qwen38KvCache::shift_scratch() { + if (!shift_scratch_.ok() && max_length_ > 1) + shift_scratch_ = DeviceTensor({max_length_ - 1, kv_dim_}, cache_dtype_, stream_); + return shift_scratch_; +} + +bool Qwen38KvCache::ok() const { + // Every group is checked: Qwen38Plugin::create relies on ok() to reject a + // state whose device allocations failed, so a partial check would report a + // broken cache as healthy and defer the failure to bind or execute time. + const auto group_ok = [this](const std::vector& group) { + if (group.size() != static_cast(num_layers_)) + return false; + for (const auto& t : group) { + if (!t.ok()) + return false; + } + return true; + }; + return group_ok(cache_k_) && group_ok(cache_v_) && group_ok(present_k_) && group_ok(present_v_); +} + +} // namespace trtmc diff --git a/src/runtime/models/qwen3_8/kv_cache.h b/src/runtime/models/qwen3_8/kv_cache.h new file mode 100644 index 000000000..434fc5ec2 --- /dev/null +++ b/src/runtime/models/qwen3_8/kv_cache.h @@ -0,0 +1,131 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +// Qwen38KvCache: autoregressive KV cache state manager. +// HF equivalent: DynamicCache / past_key_values. +// +// Manages per-layer K/V device tensors, position tracking, and attention mask +// construction. Binds directly to a TrtModule via bind_to(). + +#include "runtime/models/qwen3_8/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 Qwen38KvCacheNames { + std::vector cache_k; + std::vector cache_v; + std::vector present_k; + std::vector present_v; + std::string position_id{"position_id"}; + std::string attention_mask{"attention_mask"}; +}; + +class Qwen38KvCache : public Qwen38InferenceState { + 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. + Qwen38KvCache(int32_t num_layers, int32_t max_length, int32_t kv_dim, cudaStream_t stream, + DType cache_dtype = DType::kFloat32, Qwen38KvCacheNames names = {}); + + // --- Qwen38InferenceState 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 true; } + std::size_t device_memory_bytes() const override; + const char* state_type() const override { return "dense_kv_cache"; } + bool ok() const override; + + // --- Qwen38KvCache-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)]; } + + // Write per-layer KV produced by a batched prefill engine into the cache + // at positions [0, seq_len). Device-to-device copy on this cache's stream; + // advances position_ to seq_len. Requires seq_len <= max_length. + 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 only the cache_k/v INPUT pointers to `module`. Used for the + // prefill TrtModule whose present_k/v outputs have shape (Sq, kv_dim) + // — too big for Qwen38KvCache's single-row present buffer. The caller reads + // the prefill outputs directly from the module's own allocations and + // copies them via write_prefill_kv(). + void bind_cache_inputs(TrtModule& module); + + private: + void rebind_cache_rows(int32_t cache_rows); + std::vector mask_shape_for_engine(int32_t mask_width) 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); + + // Scratch for the cache-full row shift. cudaMemcpyAsync has undefined + // behavior on overlapping ranges, so the shift stages through this buffer. + // Allocated on first overflow, so a cache that never fills never pays for it. + DeviceTensor& shift_scratch(); + + std::vector cache_k_; // [num_layers], shape [max_length, kv_dim] + std::vector cache_v_; // [num_layers] + std::vector present_k_; // [num_layers], shape [1, kv_dim] (single step output) + std::vector present_v_; // [num_layers] + DeviceTensor shift_scratch_; // lazily sized [max_length - 1, kv_dim] + 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_; + 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)}; + Qwen38KvCacheNames names_; + TrtModule* bound_module_{nullptr}; +}; + +} // namespace trtmc diff --git a/src/runtime/models/qwen3_8/pipeline.cpp b/src/runtime/models/qwen3_8/pipeline.cpp new file mode 100644 index 000000000..6bfe3faf5 --- /dev/null +++ b/src/runtime/models/qwen3_8/pipeline.cpp @@ -0,0 +1,266 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/models/qwen3_8/pipeline.h" + +#include "runtime/models/qwen3_8/chat_templates.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { +using SteadyClock = std::chrono::steady_clock; +using TimePoint = SteadyClock::time_point; +inline double elapsed_ms(TimePoint start, TimePoint end) { + return std::chrono::duration(end - start).count(); +} +} // namespace + +namespace trtmc { + +RecurrentPipeline::RecurrentPipeline(std::unique_ptr decoder, + std::unique_ptr state, + RecurrentGenConfig config, cudaStream_t stream, + const char* name, std::shared_ptr tokenizer, + std::string model_id_str, + std::unique_ptr sampler) + : decoder_(std::move(decoder)), state_(std::move(state)), config_(config), stream_(stream), + name_(name), tokenizer_(std::move(tokenizer)), model_id_(std::move(model_id_str)), + sampler_(std::move(sampler)) { + if (!decoder_ || !decoder_->ok()) + throw std::runtime_error(std::string(name_) + ": invalid decoder module"); +} + +static std::vector encode_prompt(const ITokenizer& tokenizer, + const RecurrentGenConfig& 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 = + qwen3_8_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 RecurrentPipeline::generate(const std::string& prompt, const GenerateConfig& cfg) { + if (!tokenizer_) + throw std::runtime_error(std::string(name_) + ": 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; + int32_t eos = (cfg.eos_token_id >= 0) ? cfg.eos_token_id : config_.id_eos; + + auto sp = qwen38_sampling_params_from_config(cfg, eos); + auto output_ids = generate_from_ids(input_ids, max_new, sp); + + std::vector new_tokens( + output_ids.begin() + static_cast(input_ids.size()), output_ids.end()); + std::string text = tokenizer_->decode(new_tokens); + + return TextResult{std::move(text), std::move(new_tokens)}; +} + +RecurrentPipeline::GenerationResult +RecurrentPipeline::generate_ids(const std::vector& input_ids, const GenerateConfig& cfg) { + int32_t max_new = cfg.max_new_tokens; + int32_t eos = (cfg.eos_token_id >= 0) ? cfg.eos_token_id : config_.id_eos; + auto sp = qwen38_sampling_params_from_config(cfg, eos); + return GenerationResult{generate_from_ids(input_ids, max_new, sp)}; +} + +std::vector RecurrentPipeline::generate_from_ids(const std::vector& input_ids, + int32_t max_new_tokens, + const Qwen38SamplingParams& params) { + if (max_new_tokens == 0 || input_ids.empty()) + return input_ids; + + // Create a per-call sampler if none was injected at construction time. + Qwen38ISampler* active_sampler = sampler_.get(); + std::unique_ptr local_sampler; + if (!active_sampler) { + local_sampler = create_qwen38_sampler(params); + active_sampler = local_sampler.get(); + } + active_sampler->reset(); + + state_->reset(); + state_->bind_to(*decoder_); + + prof_prepare_ms_ = prof_forward_ms_ = prof_logits_copy_ms_ = prof_advance_ms_ = 0; + prof_steps_ = 0; + + std::vector logits; + + // ── Prefill phase ── + auto t_prefill_start = SteadyClock::now(); + for (std::size_t i = 0; i + 1 < input_ids.size(); ++i) + run_step(input_ids[i], logits); + + run_step(input_ids.back(), logits); + auto t_prefill_end = SteadyClock::now(); + + // ── Decode phase ── + std::vector output = input_ids; + const int32_t vocab_size = static_cast(logits.size()); + int32_t decode_steps = 0; + + auto t_decode_start = SteadyClock::now(); + for (int32_t step = 0; step < max_new_tokens; ++step) { + qwen38_apply_repetition_penalty(logits, params.repetition_penalty, output); + Qwen38SampleResult result = active_sampler->sample(logits.data(), vocab_size, params); + output.push_back(result.token_id); + if (result.is_eos) + break; + run_step(result.token_id, logits); + ++decode_steps; + } + auto t_decode_end = SteadyClock::now(); + + report_timing(t_prefill_start, t_prefill_end, t_decode_start, t_decode_end, + static_cast(input_ids.size()), decode_steps); + + return output; +} + +void RecurrentPipeline::report_timing(SteadyClock::time_point t_prefill_start, + SteadyClock::time_point t_prefill_end, + SteadyClock::time_point t_decode_start, + SteadyClock::time_point t_decode_end, int prefill_tokens, + int decode_steps) { + double prefill_ms = elapsed_ms(t_prefill_start, t_prefill_end); + double decode_ms = elapsed_ms(t_decode_start, t_decode_end); + double total_ms = elapsed_ms(t_prefill_start, t_decode_end); + + std::cerr << std::fixed << std::setprecision(1); + std::cerr << "[trtmc-perf] Prefill: " << prefill_tokens << " tokens, " << prefill_ms << " ms"; + if (prefill_tokens > 0) + std::cerr << " (" << std::setprecision(1) << (prefill_tokens / (prefill_ms / 1000.0)) + << " tok/s)"; + std::cerr << "\n"; + + std::cerr << "[trtmc-perf] Decode: " << decode_steps << " steps, " << decode_ms << " ms"; + if (decode_steps > 0) + std::cerr << " (" << std::setprecision(1) << (decode_steps / (decode_ms / 1000.0)) + << " tok/s, " << std::setprecision(2) << (decode_ms / decode_steps) << " ms/tok)"; + std::cerr << "\n"; + + std::cerr << "[trtmc-perf] Total generation: " << total_ms << " ms" + << " (" << (prefill_tokens + decode_steps) << " tokens)\n"; + + if (prof_steps_ > 0) { + std::cerr << std::setprecision(2); + std::cerr << "[trtmc-perf] Per-step breakdown (avg over " << prof_steps_ << " steps):\n"; + std::cerr << "[trtmc-perf] prepare_step: " << (prof_prepare_ms_ / prof_steps_) + << " ms\n"; + std::cerr << "[trtmc-perf] forward (TRT): " << (prof_forward_ms_ / prof_steps_) + << " ms\n"; + std::cerr << "[trtmc-perf] logits copy: " << (prof_logits_copy_ms_ / prof_steps_) + << " ms\n"; + std::cerr << "[trtmc-perf] state advance: " << (prof_advance_ms_ / prof_steps_) + << " ms\n"; + + std::size_t output_bytes = 0; + for (const auto& info : decoder_->output_info()) { + std::size_t n = 1; + for (auto d : info.shape) + n *= static_cast(d); + n *= dtype_size(info.dtype); + output_bytes += n; + } + std::cerr << "[trtmc-perf] D2H output size: " << std::setprecision(1) + << (output_bytes / (1024.0 * 1024.0)) << " MB (" << decoder_->output_info().size() + << " tensors)\n"; + } +} + +void RecurrentPipeline::run_step(int32_t token_id, std::vector& logits) { + auto t0 = SteadyClock::now(); + + TensorMap inputs; + + Tensor token_t; + token_t.data = &token_id; + token_t.shape = {1}; + token_t.dtype = DType::kInt32; + inputs["token_id"] = token_t; + + state_->prepare_step(inputs); + + auto t1 = SteadyClock::now(); + + // Use forward_async instead of forward() to avoid downloading + // all 63 output tensors (140+ MB of state) to CPU every step. + // Only the logits tensor (~512 KB) needs to reach CPU for sampling. + decoder_->forward_async(inputs); + + // Resolve logits device pointer + size once on first call. + if (!logits_device_ptr_) { + decoder_->sync(); // must sync before first device_ptr query + logits_device_ptr_ = decoder_->device_ptr("logits"); + if (!logits_device_ptr_) + throw std::runtime_error(std::string(name_) + ": no 'logits' output"); + for (const auto& info : decoder_->output_info()) { + if (info.name == "logits") { + logits_numel_ = 1; + for (auto d : info.shape) + logits_numel_ *= static_cast(d); + break; + } + } + if (logits_numel_ == 0) + throw std::runtime_error(std::string(name_) + ": logits tensor has zero size"); + } + + // Wait for TRT kernel to finish. + decoder_->sync(); + + auto t2 = SteadyClock::now(); + + // D2H logits only (~512 KB). Synchronous cudaMemcpy is faster than + // cudaMemcpyAsync+sync here because it bypasses stream ordering overhead. + logits.resize(logits_numel_); + // A failed copy would leave the previous step's values in `logits` and the + // sampler would silently emit a token from stale data, so surface it here. + const cudaError_t logits_copy = cudaMemcpy( + logits.data(), logits_device_ptr_, logits_numel_ * sizeof(float), cudaMemcpyDeviceToHost); + if (logits_copy != cudaSuccess) + throw std::runtime_error(std::string(name_) + ": failed to copy logits to host: " + + cudaGetErrorString(logits_copy)); + + auto t3 = SteadyClock::now(); + + // D2D state copies (present -> state) — async on the CUDA stream. + state_->advance(); + + auto t4 = SteadyClock::now(); + + prof_prepare_ms_ += elapsed_ms(t0, t1); + prof_forward_ms_ += elapsed_ms(t1, t2); + prof_logits_copy_ms_ += elapsed_ms(t2, t3); + prof_advance_ms_ += elapsed_ms(t3, t4); + ++prof_steps_; +} + +int32_t RecurrentPipeline::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/qwen3_8/pipeline.h b/src/runtime/models/qwen3_8/pipeline.h new file mode 100644 index 000000000..ac7dfdf36 --- /dev/null +++ b/src/runtime/models/qwen3_8/pipeline.h @@ -0,0 +1,89 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +// RecurrentPipeline: Qwen3.8-owned hybrid recurrent text pipeline. +// Uses Qwen38InferenceState for recurrent state ownership. + +#include "runtime/models/qwen3_8/inference_state.h" +#include "runtime/models/qwen3_8/sampler.h" +#include "trtmc/pipeline.h" +#include "trtmc/runtime/trt_module.h" +#include "trtmc/tokenizer.h" + +#include +#include +#include +#include +#include + +namespace trtmc { + +struct RecurrentGenConfig { + int32_t vocab_size{0}; + int32_t id_bos{0}; + int32_t id_eos{0}; + bool has_position_input{false}; + std::string chat_template_format{}; +}; + +class RecurrentPipeline final : public IPipeline { + public: + RecurrentPipeline(std::unique_ptr decoder, + std::unique_ptr state, RecurrentGenConfig config, + cudaStream_t stream, const char* name, + std::shared_ptr tokenizer = nullptr, + std::string model_id_str = "", + std::unique_ptr sampler = nullptr); + + 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 name_; } + + // 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); + + static int32_t argmax(const std::vector& logits); + + private: + std::unique_ptr decoder_; + std::unique_ptr state_; + RecurrentGenConfig config_; + cudaStream_t stream_; + const char* name_; + std::shared_ptr tokenizer_; + std::string model_id_; + std::unique_ptr sampler_; + + std::vector generate_from_ids(const std::vector& input_ids, + int32_t max_new_tokens, + const Qwen38SamplingParams& params); + + void run_step(int32_t token_id, std::vector& logits); + + using SteadyClock = std::chrono::steady_clock; + void report_timing(SteadyClock::time_point t_prefill_start, + SteadyClock::time_point t_prefill_end, + SteadyClock::time_point t_decode_start, SteadyClock::time_point t_decode_end, + int prefill_tokens, int decode_steps); + + // Cached logits output metadata (resolved once, reused every step) + void* logits_device_ptr_{nullptr}; + std::size_t logits_numel_{0}; + + // Per-step profiling accumulators + double prof_prepare_ms_{0}; + double prof_forward_ms_{0}; + double prof_logits_copy_ms_{0}; + double prof_advance_ms_{0}; + int prof_steps_{0}; +}; + +} // namespace trtmc diff --git a/src/runtime/models/qwen3_8/plugin.cpp b/src/runtime/models/qwen3_8/plugin.cpp new file mode 100644 index 000000000..d7a068990 --- /dev/null +++ b/src/runtime/models/qwen3_8/plugin.cpp @@ -0,0 +1,148 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +// Qwen38Plugin: handles the Qwen3.8-owned hybrid recurrent strategy. +// Qwen3.8 style models with interleaved attention and Mamba layers, +// using Qwen38KvCache for attention layers and Qwen38RecurrentState for SSM layers. + +#include "plugin_helpers.h" +#include "runtime/models/qwen3_8/hybrid_state.h" +#include "runtime/models/qwen3_8/pipeline.h" +#include "runtime/models/qwen3_8/recurrent_state.h" +#include "trtmc/runtime/distributed_runtime.h" +#include "trtmc/runtime/pipeline_registry.h" +#include "utils/json_helpers.h" + +#include +#include +#include + +namespace trtmc { + +namespace { + +struct TensorParallelRuntimeConfig { + bool enabled{false}; + int32_t tp_size{1}; +}; + +struct TensorParallelRuntime { + TensorParallelRuntimeConfig config; + DistributedRuntimeGroup group; +}; + +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); +} + +} // namespace + +class Qwen38Plugin final : public IPipelinePlugin { + public: + std::unique_ptr create(const PipelineContext& ctx) override { + load_ffi_kernels_from_bundle(ctx.bundle); + + 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); + + ModuleCreateOptions opts; + opts.runtime_cache_path = ctx.runtime_cache_path.c_str(); + opts.cuda_graphs = ctx.cuda_graphs; + if (tp_runtime.config.enabled) { + opts.distributed_communicator = tp_runtime.group.communicator; + opts.distributed_owner = tp_runtime.group.owner; + } + + const std::string engine_section = tp_runtime.config.enabled + ? tp_engine_section_name(tp_runtime.group.rank) + : std::string("engine_plan"); + auto loaded = load_trt_module_from_plan( + ctx.backend, find_section(ctx.bundle, engine_section), engine_section.c_str(), opts); + auto tokenizer = create_tokenizer_from_bundle(ctx.bundle); + + cudaStream_t stream = loaded.module->stream(); + int32_t kv_dim = compute_kv_dim(ctx.config); + + int32_t num_attention_layers = extract_json_int(ctx.config_json, "num_attention_layers", 0); + int32_t num_mamba_layers = extract_json_int(ctx.config_json, "num_mamba_layers", 0); + int32_t d_inner = extract_json_int(ctx.config_json, "d_inner", ctx.config.hidden_size * 2); + int32_t mamba_d_state = extract_json_int(ctx.config_json, "mamba_d_state", 128); + int32_t mamba_d_conv = extract_json_int(ctx.config_json, "mamba_d_conv", 4); + int32_t mamba_nheads = extract_json_int(ctx.config_json, "mamba_nheads", 0); + int32_t mamba_head_dim = extract_json_int(ctx.config_json, "mamba_head_dim", 0); + int32_t conv_dim = extract_json_int(ctx.config_json, "conv_dim", d_inner); + + // Every value above defaults to 0 when the bundle omits its key, and a + // zero silently produces a degenerate state rather than an error: no + // attention layers means an empty KV cache whose ok() is trivially + // true, and zero mamba heads means zero-element SSM tensors bound to an + // engine that expects real state. Reject them at load time instead. + const auto require_positive = [](int32_t value, const char* key) { + if (value <= 0) + throw std::runtime_error(std::string("Qwen3.8 bundle config is missing or has a " + "non-positive value for '") + + key + "'"); + }; + require_positive(num_attention_layers, "num_attention_layers"); + require_positive(num_mamba_layers, "num_mamba_layers"); + require_positive(mamba_nheads, "mamba_nheads"); + require_positive(mamba_head_dim, "mamba_head_dim"); + require_positive(mamba_d_state, "mamba_d_state"); + require_positive(mamba_d_conv, "mamba_d_conv"); + require_positive(kv_dim, "num_key_value_heads * head_dim"); + // A zero capacity would send advance() down the cache-full branch on the + // very first step, where shift_bytes is computed from max_length_ - 1. + require_positive(ctx.config.max_cache_length, "max_cache_length"); + + // Qwen38KvCache for the attention layers + DType cache_dtype = cache_dtype_from_precision(ctx.config.precision); + auto cache = std::make_unique( + num_attention_layers, ctx.config.max_cache_length, kv_dim, stream, cache_dtype); + if (!cache->ok()) + throw std::runtime_error("Failed to create Qwen38KvCache for hybrid model"); + + // Qwen38RecurrentState for the Mamba/SSM layers (conv_state + ssm_state) + int32_t effective_conv_dim = (conv_dim > 0) ? conv_dim : d_inner; + int64_t conv_elems = static_cast(effective_conv_dim) * mamba_d_conv; + int64_t ssm_elems = + static_cast(mamba_nheads) * std::max(mamba_head_dim, 1) * mamba_d_state; + + auto ssm = + std::make_unique(num_mamba_layers, + std::vector{ + {"conv_state", {conv_elems}, "present_conv"}, + {"ssm_state", {ssm_elems}, "present_ssm"}}, + stream); + + if (!ssm->ok()) + throw std::runtime_error("Failed to create Qwen38RecurrentState for hybrid model"); + + auto hybrid = std::make_unique(std::move(cache), std::move(ssm)); + if (!hybrid->ok()) + throw std::runtime_error("Failed to create Qwen38HybridState for hybrid model"); + auto rgc = make_recurrent_gen_config(ctx.config); + rgc.has_position_input = loaded.module->has_input("position_id"); + apply_recurrent_chat_template_format(ctx.bundle, rgc); + + return std::make_unique(std::move(loaded.module), std::move(hybrid), rgc, + stream, "HybridPipeline", std::move(tokenizer), + ctx.bundle.info.model_id); + } +}; + +REGISTER_PIPELINE_PLUGIN_WITH_MANIFEST(register_qwen3_8_plugin, Qwen38Plugin, + "qwen3_8_hybrid_mamba_attention"); + +} // namespace trtmc diff --git a/src/runtime/models/qwen3_8/plugin_helpers.cpp b/src/runtime/models/qwen3_8/plugin_helpers.cpp new file mode 100644 index 000000000..b463f407d --- /dev/null +++ b/src/runtime/models/qwen3_8/plugin_helpers.cpp @@ -0,0 +1,551 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "plugin_helpers.h" + +#include "runtime/models/qwen3_8/chat_templates.h" +#include "trtmc/runtime/trt_backend.h" +#include "utils/json_helpers.h" + +#include +#include +#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); + +} // namespace + +void log_trt_load_timing(const char* label, double load_deserialize_ms, std::size_t plan_bytes) { + std::ostringstream line; + line << std::fixed << std::setprecision(6) << "[trtmc.load_timing] label=\"" + << (label ? label : "engine") << "\" load_deserialize_ms=" << load_deserialize_ms + << " plan_bytes=" << plan_bytes; + std::cerr << line.str() << '\n'; +} + +// Tokenizer helpers. + +bool detect_add_special_tokens(const BundleFile& bundle) { + if (bundle.info.tokenizer_add_special_tokens_present) + return bundle.info.tokenizer_add_special_tokens; + + auto* config_data = find_section(bundle, "config.json"); + if (!config_data) + return true; + std::string cfg_text(config_data->begin(), config_data->end()); + auto pos = cfg_text.find("\"tokenizer_add_special_tokens\""); + if (pos == std::string::npos) + return true; + auto val_pos = cfg_text.find(':', pos); + if (val_pos == std::string::npos) + return true; + auto value_pos = cfg_text.find_first_not_of(" \t\r\n", val_pos + 1); + if (value_pos == std::string::npos) + return true; + if (cfg_text.compare(value_pos, 5, "false") == 0 || cfg_text[value_pos] == '0') + return false; + if (cfg_text.compare(value_pos, 4, "true") == 0 || cfg_text[value_pos] == '1') + return true; + return true; +} + +namespace { + +TokenizerSpecialFrame detect_tokenizer_special_frame(const BundleFile& bundle) { + TokenizerSpecialFrame frame; + auto* config_data = find_section(bundle, "config.json"); + if (!config_data) + return frame; + std::string cfg_text(config_data->begin(), config_data->end()); + const bool has_prefix = cfg_text.find("\"tokenizer_special_prefix_ids\"") != std::string::npos; + const bool has_suffix = cfg_text.find("\"tokenizer_special_suffix_ids\"") != std::string::npos; + if (!has_prefix && !has_suffix) + return frame; + + frame.present = true; + frame.prefix = extract_json_int_array(cfg_text, "tokenizer_special_prefix_ids"); + frame.suffix = extract_json_int_array(cfg_text, "tokenizer_special_suffix_ids"); + return frame; +} + +std::shared_ptr apply_tokenizer_special_frame(std::unique_ptr tokenizer, + const TokenizerSpecialFrame& frame) { + if (!tokenizer) + return nullptr; + std::shared_ptr shared(std::move(tokenizer)); + if (!frame.present || (frame.prefix.empty() && frame.suffix.empty())) + return shared; + return std::make_shared(std::move(shared), frame.prefix, frame.suffix); +} + +TokenizerSpecialFrame detect_requested_tokenizer_special_frame(const BundleFile& bundle, + bool add_special_tokens) { + if (!add_special_tokens) + return TokenizerSpecialFrame{}; + return detect_tokenizer_special_frame(bundle); +} + +std::shared_ptr try_create_native_tokenizer_kind(TokenizerFactory factory, + const char* data, std::size_t size, + bool add_special_tokens, + const TokenizerSpecialFrame& frame, + const char* label) { + try { + auto tok = factory(data, size, add_special_tokens); + if (!tok) + return nullptr; + std::cerr << "[trtmc] Using native " << label << " tokenizer" << std::endl; + return apply_tokenizer_special_frame(std::move(tok), frame); + } catch (...) { + return nullptr; + } +} + +} // namespace + +bool is_bpe_tokenizer_json(const BundleFile& bundle) { + auto* tok_data = find_section(bundle, "tokenizer.json"); + if (!tok_data || tok_data->empty()) + return false; + // Quick string search — avoid full JSON parse just for type detection + std::string_view json(tok_data->data(), tok_data->size()); + return json.find("\"type\":\"BPE\"") != std::string_view::npos || + json.find("\"type\": \"BPE\"") != std::string_view::npos; +} + +std::shared_ptr try_create_native_bpe(const BundleFile& bundle, bool add_special, + bool throw_on_failure) { + auto* tok_data = find_section(bundle, "tokenizer.json"); + if (!tok_data || tok_data->empty()) + return nullptr; + try { + auto tok = CreateBpeTokenizer(tok_data->data(), tok_data->size(), add_special); + if (tok) { + std::cerr << "[trtmc] Using native BPE tokenizer" << std::endl; + } + return tok; + } catch (const std::exception& e) { + // "Not a BPE tokenizer" -> non-BPE model (WordPiece, Unigram), allow fallback + std::string msg = e.what(); + bool is_non_bpe = msg.find("Not a BPE") != std::string::npos; + + if (throw_on_failure || (!is_non_bpe && is_bpe_tokenizer_json(bundle))) { + // BPE model but native failed -> error, no silent fallback + throw std::runtime_error(std::string("Native BPE tokenizer failed for BPE model: ") + + e.what()); + } + std::cerr << "[trtmc] Native BPE unavailable (" << e.what() + << "), falling back to HF Python" << std::endl; + } + return nullptr; +} + +std::shared_ptr try_create_native_tokenizer(const BundleFile& bundle, + bool add_special_tokens) { + auto* tok_data = find_section(bundle, "tokenizer.json"); + if (!tok_data || tok_data->empty()) + return nullptr; + + const char* data = tok_data->data(); + std::size_t size = tok_data->size(); + const auto special_frame = detect_requested_tokenizer_special_frame(bundle, add_special_tokens); + const bool native_add_special = !special_frame.present && add_special_tokens; + + if (auto tokenizer = try_create_native_tokenizer_kind(CreateBpeTokenizer, data, size, + native_add_special, special_frame, "BPE")) + return tokenizer; + + if (auto tokenizer = try_create_native_tokenizer_kind( + CreateWordPieceTokenizer, data, size, native_add_special, special_frame, "WordPiece")) + return tokenizer; + + return try_create_native_tokenizer_kind(CreateUnigramTokenizer, data, size, native_add_special, + special_frame, "Unigram"); +} + +std::shared_ptr create_tokenizer_from_bundle(const BundleFile& bundle) { + bool add_special = detect_add_special_tokens(bundle); + return try_create_native_tokenizer(bundle, add_special); +} + +// TRT module loading (delegated to IBackend). + +LoadedModule load_trt_module_from_plan(IBackend* backend, const std::vector* plan, + const char* label, const ModuleCreateOptions& options) { + if (!plan || plan->empty()) + throw std::runtime_error(std::string("Bundle missing ") + label); + if (!backend) + throw std::runtime_error("No backend loaded"); + + LoadedModule result; + const auto t0 = SteadyClock::now(); + result.module = backend->create_module(plan->data(), plan->size(), options); + const auto t1 = SteadyClock::now(); + log_trt_load_timing(label, elapsed_ms(t0, t1), plan->size()); + if (!result.module || !result.module->ok()) + throw std::runtime_error(std::string("Failed to create ITrtModule for ") + label); + result.module->set_timing_label(label ? label : "engine"); + return result; +} + +LoadedModule try_load_trt_module_from_plan(IBackend* backend, const std::vector* plan, + const char* label, const ModuleCreateOptions& options) { + if (!plan || plan->empty()) + return LoadedModule{}; + try { + return load_trt_module_from_plan(backend, plan, label, options); + } catch (...) { + std::cerr << "[trtmc] WARNING: failed to load optional engine: " << label << std::endl; + return LoadedModule{}; + } +} + +std::unique_ptr extract_optional_module(IBackend* backend, + const std::vector* plan, + const char* label, + const ModuleCreateOptions& options) { + auto loaded = try_load_trt_module_from_plan(backend, plan, label, options); + if (loaded.module && loaded.module->ok()) + return std::move(loaded.module); + return nullptr; +} + +// Dual-profile module loading (delegated to IBackend). + +DualProfileModules load_dual_profile_modules(IBackend* backend, const std::vector* plan, + const char* label, + const ModuleCreateOptions& options) { + if (!plan || plan->empty()) + throw std::runtime_error(std::string("Bundle missing ") + label); + if (!backend) + throw std::runtime_error("No backend loaded"); + + const auto t0 = SteadyClock::now(); + auto pair = backend->create_dual_profile_modules(plan->data(), plan->size(), options); + const auto t1 = SteadyClock::now(); + log_trt_load_timing(label, elapsed_ms(t0, t1), plan->size()); + if (!pair.decode || !pair.decode->ok()) + throw std::runtime_error(std::string("Failed to create dual-profile modules for ") + label); + + DualProfileModules out; + out.prefill = std::move(pair.prefill); + out.decode = std::move(pair.decode); + if (out.prefill) + out.prefill->set_timing_label(std::string(label ? label : "engine") + ":prefill"); + if (out.decode) + out.decode->set_timing_label(std::string(label ? label : "engine") + ":decode"); + return out; +} + +// Config helpers. + +int32_t compute_kv_dim(const BaseConfig& cfg) { + int32_t hd = (cfg.head_dim > 0) ? cfg.head_dim + : ((cfg.num_heads > 0) ? cfg.hidden_size / cfg.num_heads : 128); + int32_t kv_heads = (cfg.num_kv_heads > 0) ? cfg.num_kv_heads : cfg.num_heads; + return kv_heads * hd; +} + +DType cache_dtype_from_precision(const std::string& precision) { + if (precision == "fp16") + return DType::kFloat16; + if (precision == "bf16") + return DType::kBFloat16; + return DType::kFloat32; +} + +RecurrentGenConfig make_recurrent_gen_config(const BaseConfig& cfg) { + RecurrentGenConfig rgc; + rgc.vocab_size = cfg.vocab_size; + rgc.id_bos = cfg.id_bos; + rgc.id_eos = cfg.id_eos; + return rgc; +} + +void apply_recurrent_chat_template_format(const BundleFile& bundle, RecurrentGenConfig& rgc) { + 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()); + } + rgc.chat_template_format = qwen3_8_detect_chat_template_format(chat_tpl); +} + +// 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 = '_'; + } + // mkstemp in a private 0700 directory, not a predictable /tmp path: the + // old name was world-writable and could be pre-created as a symlink or + // swapped between the write and the dlopen, which would hand this process + // attacker-controlled code. The write status is checked so a short write + // never reaches the loader. + std::string dir_template = "/tmp/trtmc_kernels_XXXXXX"; + std::vector dir_buf(dir_template.begin(), dir_template.end()); + dir_buf.push_back('\0'); + if (mkdtemp(dir_buf.data()) == nullptr) + return {}; + + std::string tmp_path = std::string(dir_buf.data()) + "/" + safe_name + ".so"; + const int fd = + ::open(tmp_path.c_str(), O_WRONLY | O_CREAT | O_EXCL | O_NOFOLLOW, S_IRUSR | S_IWUSR); + if (fd < 0) { + ::rmdir(dir_buf.data()); + return {}; + } + + std::size_t written = 0; + bool write_ok = true; + while (written < size) { + const ssize_t n = ::write(fd, data + written, size - written); + if (n <= 0) { + write_ok = false; + break; + } + written += static_cast(n); + } + if (::close(fd) != 0) + write_ok = false; + if (!write_ok) { + ::unlink(tmp_path.c_str()); + ::rmdir(dir_buf.data()); + return {}; + } + return tmp_path; +} + +// Remove the staged .so and its private directory once the loader is done. +void remove_kernel_so_temp(const std::string& tmp_path) { + if (tmp_path.empty()) + return; + ::unlink(tmp_path.c_str()); + const auto slash = tmp_path.find_last_of('/'); + if (slash != std::string::npos) + ::rmdir(tmp_path.substr(0, slash).c_str()); +} + +// 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 (tmp_path.empty()) { + std::cerr << "[ffi] Failed to stage kernel .so for: " << global_name << '\n'; + return; + } + 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'; + } + remove_kernel_so_temp(tmp_path); +} + +// 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/qwen3_8/plugin_helpers.h b/src/runtime/models/qwen3_8/plugin_helpers.h new file mode 100644 index 000000000..9adf9fb70 --- /dev/null +++ b/src/runtime/models/qwen3_8/plugin_helpers.h @@ -0,0 +1,139 @@ +/* + * 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/qwen3_8/inference_state.h" +#include "runtime/models/qwen3_8/kv_cache.h" +#include "runtime/models/qwen3_8/pipeline.h" +#include "trtmc/runtime/pipeline_plugin.h" +#include "trtmc/runtime/trt_backend.h" +#include "trtmc/runtime/trt_module.h" +#include "trtmc/tokenizer.h" + +#include +#include +#include +#include + +namespace trtmc { + +// A loaded TRT engine, ready for inference. +// The stream is owned internally by the module — callers get it via module->stream(). +struct LoadedModule { + std::unique_ptr module; +}; + +// Load a TRT engine from a serialized plan via the backend. Throws on failure. +LoadedModule load_trt_module_from_plan(IBackend* backend, const std::vector* plan, + const char* label, const ModuleCreateOptions& options = {}); + +// Emit a parseable runtime load/deserialization timing line. +void log_trt_load_timing(const char* label, double load_deserialize_ms, std::size_t plan_bytes); + +// Like load_trt_module_from_plan but returns empty LoadedModule on failure +// instead of throwing (for optional engines). +LoadedModule try_load_trt_module_from_plan(IBackend* backend, const std::vector* plan, + const char* label, + const ModuleCreateOptions& options = {}); + +// Load an optional TRT module, returning nullptr if the plan is absent. +// On deserialization failure, returns nullptr (does not throw). +std::unique_ptr extract_optional_module(IBackend* backend, + const std::vector* plan, + const char* label, + const ModuleCreateOptions& options = {}); + +// Dual-profile TRT module group: one shared backend engine, two module +// contexts (one per optimization profile). Weights live once in GPU memory +// and both modules share the CUDA stream. Use `decode->stream()` to obtain +// the shared stream. +struct DualProfileModules { + std::unique_ptr prefill; // batched Sq profile (null if single-profile) + std::unique_ptr decode; // Sq=1 profile, or the only profile if single-profile +}; + +// Load an engine from a serialized plan via the backend and create two +// execution contexts — one per optimization profile — sharing the engine. +// When the engine has fewer than 2 profiles, `prefill` is left null and +// `decode` holds the single-profile context (legacy bundles). +DualProfileModules load_dual_profile_modules(IBackend* backend, const std::vector* plan, + const char* label, + const ModuleCreateOptions& options = {}); + +// Detect whether the bundle's config requests add_special_tokens for the tokenizer. +bool detect_add_special_tokens(const BundleFile& bundle); + +// Check if the bundle's tokenizer.json describes a BPE model. +bool is_bpe_tokenizer_json(const BundleFile& bundle); + +// Try to create a native C++ BPE tokenizer from the bundle's tokenizer.json. +// Returns nullptr if the section is absent or the model is non-BPE. +// If throw_on_failure is true, throws instead of returning nullptr on BPE parse errors. +std::shared_ptr try_create_native_bpe(const BundleFile& bundle, bool add_special, + bool throw_on_failure); + +// Try to create a native C++ tokenizer from the bundle's tokenizer.json. +// Attempts: BPE -> WordPiece -> Unigram. Returns nullptr if none match. +std::shared_ptr try_create_native_tokenizer(const BundleFile& bundle, + bool add_special_tokens); + +// Create a native tokenizer from bundle. Tries BPE -> WordPiece -> Unigram. +// Returns nullptr if no native tokenizer matches. +std::shared_ptr create_tokenizer_from_bundle(const BundleFile& bundle); + +// Compute the KV cache dimension from model config. +int32_t compute_kv_dim(const BaseConfig& cfg); + +// Convert the BaseConfig precision string ("fp16", "bf16", "fp32") to a DType +// for use as KV cache element type. +DType cache_dtype_from_precision(const std::string& precision); + +// Build a RecurrentGenConfig from model config fields. +RecurrentGenConfig make_recurrent_gen_config(const BaseConfig& cfg); + +// Populate recurrent chat-template metadata from tokenizer_config.json. +void apply_recurrent_chat_template_format(const BundleFile& bundle, RecurrentGenConfig& rgc); + +// 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/qwen3_8/qwen3_8_recurrent_step_contracts.h b/src/runtime/models/qwen3_8/qwen3_8_recurrent_step_contracts.h new file mode 100644 index 000000000..49da6a0b6 --- /dev/null +++ b/src/runtime/models/qwen3_8/qwen3_8_recurrent_step_contracts.h @@ -0,0 +1,65 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include +#include + +namespace trtmc { +namespace qwen3_8_recurrent { + +// A null entry, or a negative expected count, is a contract violation rather +// than a pass: these validators run before the recurrent step reads the +// vectors, so they must not dereference what they are asked to check. +template +bool validate_state_layer_count(const std::array>*, N>& states, + int32_t expected_layers) { + if (expected_layers < 0) { + return false; + } + for (const auto* state : states) { + if (state == nullptr || static_cast(state->size()) != expected_layers) { + return false; + } + } + return true; +} + +struct StateTensorView { + const std::vector>* values{nullptr}; + std::size_t expected_elems{0}; +}; + +template +bool validate_state_tensor_sizes(const std::array& states, int32_t num_layers) { + if (num_layers < 0) { + return false; + } + const auto expected = static_cast(num_layers); + for (const auto& state : states) { + if (state.values == nullptr || state.values->size() < expected) { + return false; + } + } + for (std::size_t idx = 0; idx < expected; ++idx) { + for (const auto& state : states) { + if ((*state.values)[idx].size() != state.expected_elems) { + return false; + } + } + } + return true; +} + +inline void initialize_layer_outputs(int32_t num_layers, std::size_t elems, + std::vector>& outputs) { + outputs.assign(static_cast(num_layers), std::vector(elems, 0.0F)); +} + +} // namespace qwen3_8_recurrent +} // namespace trtmc diff --git a/src/runtime/models/qwen3_8/recurrent_output_initializers.h b/src/runtime/models/qwen3_8/recurrent_output_initializers.h new file mode 100644 index 000000000..e88729622 --- /dev/null +++ b/src/runtime/models/qwen3_8/recurrent_output_initializers.h @@ -0,0 +1,42 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "runtime/models/qwen3_8/qwen3_8_recurrent_step_contracts.h" + +#include +#include +#include + +namespace trtmc { +namespace qwen3_8_recurrent { + +inline void initialize_rwkv_outputs(int32_t num_layers, int32_t vocab_size, std::size_t state_elems, + std::vector& logits, + std::vector>& present_attn_by_layer, + std::vector>& present_ff_by_layer, + std::vector>& present_num_by_layer, + std::vector>& present_den_by_layer, + std::vector>& present_max_by_layer) { + logits.assign(static_cast(vocab_size), 0.0F); + initialize_layer_outputs(num_layers, state_elems, present_attn_by_layer); + initialize_layer_outputs(num_layers, state_elems, present_ff_by_layer); + initialize_layer_outputs(num_layers, state_elems, present_num_by_layer); + initialize_layer_outputs(num_layers, state_elems, present_den_by_layer); + initialize_layer_outputs(num_layers, state_elems, present_max_by_layer); +} + +inline void initialize_mamba_outputs(int32_t num_layers, int32_t vocab_size, std::size_t conv_elems, + std::size_t ssm_elems, std::vector& logits, + std::vector>& present_conv_by_layer, + std::vector>& present_ssm_by_layer) { + logits.assign(static_cast(vocab_size), 0.0F); + initialize_layer_outputs(num_layers, conv_elems, present_conv_by_layer); + initialize_layer_outputs(num_layers, ssm_elems, present_ssm_by_layer); +} + +} // namespace qwen3_8_recurrent +} // namespace trtmc diff --git a/src/runtime/models/qwen3_8/recurrent_state.cpp b/src/runtime/models/qwen3_8/recurrent_state.cpp new file mode 100644 index 000000000..f60544d55 --- /dev/null +++ b/src/runtime/models/qwen3_8/recurrent_state.cpp @@ -0,0 +1,103 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/models/qwen3_8/recurrent_state.h" + +#include "trtmc/runtime/trt_module.h" + +#include +#include +#include + +namespace trtmc { + +Qwen38RecurrentState::Qwen38RecurrentState(int32_t num_layers, std::vector specs, + cudaStream_t stream) + : specs_(std::move(specs)), num_layers_(num_layers), stream_(stream) { + state_.resize(specs_.size()); + present_.resize(specs_.size()); + + for (std::size_t si = 0; si < specs_.size(); ++si) { + state_[si].reserve(static_cast(num_layers)); + present_[si].reserve(static_cast(num_layers)); + for (int32_t li = 0; li < num_layers; ++li) { + state_[si].emplace_back(specs_[si].shape, DType::kFloat32, stream); + present_[si].emplace_back(specs_[si].shape, DType::kFloat32, stream); + } + } + + reset(); +} + +void Qwen38RecurrentState::bind_to(TrtModule& module) { + for (std::size_t si = 0; si < specs_.size(); ++si) { + const auto& name = specs_[si].name; + const auto& out_prefix = + specs_[si].output_prefix.empty() ? ("present_" + name) : specs_[si].output_prefix; + for (int32_t li = 0; li < num_layers_; ++li) { + auto suffix = "_" + std::to_string(li); + auto uli = static_cast(li); + + module.bind_external(name + suffix, state_[si][uli].data()); + module.bind_external(out_prefix + suffix, present_[si][uli].data()); + } + } +} + +void Qwen38RecurrentState::prepare_step(TensorMap& /*inputs*/, int32_t /*seq_len*/) {} + +void Qwen38RecurrentState::advance(int32_t n_tokens) { + for (std::size_t si = 0; si < specs_.size(); ++si) { + for (int32_t li = 0; li < num_layers_; ++li) { + auto uli = static_cast(li); + state_[si][uli].copy_from(present_[si][uli]); + } + } + position_ += n_tokens; +} + +void Qwen38RecurrentState::reset() { + position_ = 0; + for (std::size_t si = 0; si < specs_.size(); ++si) { + for (int32_t li = 0; li < num_layers_; ++li) { + auto uli = static_cast(li); + cudaMemsetAsync(state_[si][uli].data(), 0, state_[si][uli].nbytes(), stream_); + cudaMemsetAsync(present_[si][uli].data(), 0, present_[si][uli].nbytes(), stream_); + } + } + cudaStreamSynchronize(stream_); +} + +std::size_t Qwen38RecurrentState::device_memory_bytes() const { + std::size_t total = 0; + for (std::size_t si = 0; si < specs_.size(); ++si) { + for (const auto& t : state_[si]) + total += t.nbytes(); + for (const auto& t : present_[si]) + total += t.nbytes(); + } + return total; +} + +bool Qwen38RecurrentState::ok() const { + // Both groups are checked: a present_ allocation failure is just as fatal + // as a state_ one, and Qwen38Plugin::create rejects the pipeline on ok(). + const auto group_ok = [this](const std::vector& group) { + if (group.size() != static_cast(num_layers_)) + return false; + for (const auto& t : group) { + if (!t.ok()) + return false; + } + return true; + }; + for (std::size_t si = 0; si < specs_.size(); ++si) { + if (!group_ok(state_[si]) || !group_ok(present_[si])) + return false; + } + return true; +} + +} // namespace trtmc diff --git a/src/runtime/models/qwen3_8/recurrent_state.h b/src/runtime/models/qwen3_8/recurrent_state.h new file mode 100644 index 000000000..3d0b806df --- /dev/null +++ b/src/runtime/models/qwen3_8/recurrent_state.h @@ -0,0 +1,51 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "runtime/models/qwen3_8/inference_state.h" +#include "trtmc/runtime/device_tensor.h" + +#include +#include +#include +#include + +namespace trtmc { + +class Qwen38RecurrentState final : public Qwen38InferenceState { + public: + struct TensorSpec { + std::string name; + std::vector shape; + std::string output_prefix; + }; + + Qwen38RecurrentState(int32_t num_layers, std::vector specs, cudaStream_t stream); + + 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 -1; } + int32_t num_layers() const override { return num_layers_; } + bool needs_attention_mask() const override { return false; } + std::size_t device_memory_bytes() const override; + const char* state_type() const override { return "qwen3_8_recurrent"; } + bool ok() const override; + + const std::vector& specs() const { return specs_; } + + private: + std::vector specs_; + std::vector> state_; + std::vector> present_; + int32_t num_layers_{0}; + int32_t position_{0}; + cudaStream_t stream_{nullptr}; +}; + +} // namespace trtmc diff --git a/src/runtime/models/qwen3_8/sampler.cpp b/src/runtime/models/qwen3_8/sampler.cpp new file mode 100644 index 000000000..5d185d84d --- /dev/null +++ b/src/runtime/models/qwen3_8/sampler.cpp @@ -0,0 +1,276 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/models/qwen3_8/sampler.h" + +#include "trtmc/pipeline.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace trtmc { + +namespace { + +Qwen38SampleResult argmax_over_logits(const float* logits, int32_t vocab_size, + int32_t eos_token_id) { + Qwen38SampleResult 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 = (result.token_id == eos_token_id); + return result; +} + +struct FilteredDistribution { + std::vector indices; + std::vector probs; + int32_t keep{0}; +}; + +constexpr float kSamplingEpsilon = 1e-6F; + +float sanitized_temperature(float temperature) { + if (!std::isfinite(temperature)) + return 1.0F; + return std::max(temperature, 0.0F); +} + +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); +} + +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); +} + +bool top_p_enabled(float top_p) { + return top_p > 0.0F && top_p < 1.0F - kSamplingEpsilon; +} + +bool greedy_equivalent(const Qwen38SamplingParams& 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; +} + +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; + } +} + +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); +} + +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); +} + +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; + } +} + +FilteredDistribution build_filtered_distribution(const float* logits, int32_t vocab_size, + const Qwen38SamplingParams& 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; +} + +class Qwen38GreedySampler final : public Qwen38ISampler { + public: + Qwen38SampleResult sample(const float* logits, int32_t vocab_size, + const Qwen38SamplingParams& params) override { + if (vocab_size <= 0 || logits == nullptr) { + Qwen38SampleResult result; + result.token_id = 0; + result.is_eos = (0 == params.eos_token_id); + return result; + } + return argmax_over_logits(logits, vocab_size, params.eos_token_id); + } + + Qwen38LogitsLocation logits_location() const override { return Qwen38LogitsLocation::HOST; } + const char* sampler_type() const override { return "qwen38_greedy"; } +}; + +class Qwen38TopKSampler final : public Qwen38ISampler { + public: + explicit Qwen38TopKSampler(uint64_t initial_seed) + : rng_state_(initial_seed == 0 ? 1 : initial_seed), + initial_seed_(initial_seed == 0 ? 1 : initial_seed) {} + + Qwen38SampleResult sample(const float* logits, int32_t vocab_size, + const Qwen38SamplingParams& params) override { + Qwen38SampleResult result; + if (vocab_size <= 0 || logits == nullptr) { + result.token_id = 0; + result.is_eos = (0 == params.eos_token_id); + return result; + } + + if (greedy_equivalent(params)) + return argmax_over_logits(logits, vocab_size, params.eos_token_id); + + const FilteredDistribution dist = build_filtered_distribution(logits, vocab_size, params); + + rng_state_ ^= rng_state_ << 13; + rng_state_ ^= rng_state_ >> 7; + rng_state_ ^= rng_state_ << 17; + const float u = static_cast(rng_state_ & 0xFFFFFFFF) / 4294967296.0F; + + 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 = (result.token_id == params.eos_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 = (result.token_id == params.eos_token_id); + return result; + } + + Qwen38LogitsLocation logits_location() const override { return Qwen38LogitsLocation::HOST; } + const char* sampler_type() const override { return "qwen38_top_k"; } + void reset() override { rng_state_ = initial_seed_; } + + private: + uint64_t rng_state_; + uint64_t initial_seed_; +}; + +} // namespace + +Qwen38SamplingParams qwen38_sampling_params_from_config(const GenerateConfig& cfg, + int32_t default_eos) { + Qwen38SamplingParams p; + p.temperature = cfg.temperature; + p.top_k = cfg.top_k; + p.top_p = cfg.top_p; + p.min_p = cfg.min_p; + p.repetition_penalty = cfg.repetition_penalty; + p.seed = cfg.seed; + p.eos_token_id = (cfg.eos_token_id >= 0) ? cfg.eos_token_id : default_eos; + return p; +} + +void qwen38_apply_repetition_penalty(std::vector& logits, float penalty, + const std::vector& token_history) { + if (penalty <= 0.0F || std::fabs(penalty - 1.0F) < kSamplingEpsilon || logits.empty()) + return; + const auto vocab_size = static_cast(logits.size()); + std::unordered_set seen; + seen.reserve(token_history.size()); + for (int32_t token : token_history) { + if (token < 0 || token >= vocab_size) + continue; + if (!seen.insert(token).second) + continue; + float& score = logits[static_cast(token)]; + score = score < 0.0F ? score * penalty : score / penalty; + } +} + +std::unique_ptr create_qwen38_sampler(const Qwen38SamplingParams& params) { + 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(); + } + + const uint64_t seed = (params.seed >= 0) ? static_cast(params.seed) : 42ULL; + return std::make_unique(seed); +} + +} // namespace trtmc diff --git a/src/runtime/models/qwen3_8/sampler.h b/src/runtime/models/qwen3_8/sampler.h new file mode 100644 index 000000000..5783a3db9 --- /dev/null +++ b/src/runtime/models/qwen3_8/sampler.h @@ -0,0 +1,58 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +namespace trtmc { + +struct GenerateConfig; + +struct Qwen38SamplingParams { + float temperature{1.0F}; + int32_t top_k{1}; + float top_p{1.0F}; + float min_p{0.0F}; + float repetition_penalty{1.0F}; + int32_t seed{-1}; + int32_t eos_token_id{-1}; +}; + +enum class Qwen38LogitsLocation { + HOST, + DEVICE, +}; + +struct Qwen38SampleResult { + int32_t token_id{0}; + float logprob{0.0F}; + bool is_eos{false}; +}; + +class Qwen38ISampler { + public: + virtual ~Qwen38ISampler() = default; + virtual Qwen38SampleResult sample(const float* logits, int32_t vocab_size, + const Qwen38SamplingParams& params) = 0; + virtual Qwen38LogitsLocation logits_location() const = 0; + virtual const char* sampler_type() const = 0; + virtual void reset() {} +}; + +Qwen38SamplingParams qwen38_sampling_params_from_config(const GenerateConfig& cfg, + int32_t default_eos = -1); + +// Scale down logits for tokens already present in token_history. Applied to the +// logits before sampling, so both samplers honor it without either needing to +// know the history. A penalty of 1.0 (the default) leaves logits untouched. +void qwen38_apply_repetition_penalty(std::vector& logits, float penalty, + const std::vector& token_history); + +std::unique_ptr create_qwen38_sampler(const Qwen38SamplingParams& params); + +} // namespace trtmc diff --git a/tests/cpp/models/qwen3_8/test_qwen3_8_recurrent_output_initializers.cpp b/tests/cpp/models/qwen3_8/test_qwen3_8_recurrent_output_initializers.cpp new file mode 100644 index 000000000..5286006ac --- /dev/null +++ b/tests/cpp/models/qwen3_8/test_qwen3_8_recurrent_output_initializers.cpp @@ -0,0 +1,102 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +// Unit tests for recurrent-owned output initializers. + +#include "runtime/models/qwen3_8/recurrent_output_initializers.h" + +#include +#include +#include + +namespace { + +namespace under_test = trtmc::qwen3_8_recurrent; + +int g_failures = 0; + +void check(bool condition, const char* name) { + if (!condition) { + std::cerr << "FAIL: " << name << '\n'; + ++g_failures; + } +} + +void test_initialize_rwkv_outputs() { + std::vector logits; + std::vector> attn; + std::vector> ff; + std::vector> num; + std::vector> den; + std::vector> maxv; + + under_test::initialize_rwkv_outputs(3, 11, 5, logits, attn, ff, num, den, maxv); + + check(logits.size() == 11, "rwkv outputs allocate logits"); + check(attn.size() == 3 && attn[0].size() == 5, "rwkv outputs allocate attn"); + check(ff.size() == 3 && ff[1].size() == 5, "rwkv outputs allocate ff"); + check(num.size() == 3 && num[2].size() == 5, "rwkv outputs allocate num"); + check(den.size() == 3 && den[0].size() == 5, "rwkv outputs allocate den"); + check(maxv.size() == 3 && maxv[1].size() == 5, "rwkv outputs allocate max"); +} + +void test_initialize_mamba_outputs() { + std::vector logits; + std::vector> conv; + std::vector> ssm; + + under_test::initialize_mamba_outputs(2, 13, 6, 7, logits, conv, ssm); + + check(logits.size() == 13, "mamba outputs allocate logits"); + check(conv.size() == 2 && conv[0].size() == 6, "mamba outputs allocate conv"); + check(ssm.size() == 2 && ssm[1].size() == 7, "mamba outputs allocate ssm"); +} + +void test_qwen3_8_recurrent_contracts() { + std::vector> a(2, std::vector(4, 1.0F)); + std::vector> b(2, std::vector(4, 2.0F)); + std::vector> bad_layers(1, std::vector(4, 0.0F)); + std::vector> bad_sizes = a; + bad_sizes[1].resize(3); + + const auto ok_states = std::array>*, 2>{&a, &b}; + const auto bad_layer_states = + std::array>*, 2>{&a, &bad_layers}; + + check(under_test::validate_state_layer_count(ok_states, 2), + "qwen3.8 contract accepts matching layer count"); + check(!under_test::validate_state_layer_count(bad_layer_states, 2), + "qwen3.8 contract rejects layer count mismatch"); + + const auto ok_specs = std::array{ + under_test::StateTensorView{&a, 4}, under_test::StateTensorView{&b, 4}}; + const auto bad_specs = std::array{ + under_test::StateTensorView{&a, 4}, under_test::StateTensorView{&bad_sizes, 4}}; + + check(under_test::validate_state_tensor_sizes(ok_specs, 2), + "qwen3.8 contract accepts matching tensor sizes"); + check(!under_test::validate_state_tensor_sizes(bad_specs, 2), + "qwen3.8 contract rejects tensor size mismatch"); + + std::vector> outputs; + under_test::initialize_layer_outputs(3, 2, outputs); + check(outputs.size() == 3, "qwen3.8 contract allocates layer outputs"); + check(outputs[0] == std::vector({0.0F, 0.0F}), + "qwen3.8 contract initializes outputs to zero"); +} + +} // namespace + +int main() { + test_initialize_rwkv_outputs(); + test_initialize_mamba_outputs(); + test_qwen3_8_recurrent_contracts(); + + if (g_failures != 0) { + std::cerr << g_failures << " recurrent output initializer test(s) failed\n"; + return 1; + } + return 0; +} diff --git a/tests/cpp/models/qwen3_8/test_qwen3_8_recurrent_pipeline.cpp b/tests/cpp/models/qwen3_8/test_qwen3_8_recurrent_pipeline.cpp new file mode 100644 index 000000000..197e7c09f --- /dev/null +++ b/tests/cpp/models/qwen3_8/test_qwen3_8_recurrent_pipeline.cpp @@ -0,0 +1,337 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +// ============================================================================= +// ISO 26262 Traceability +// ============================================================================= +// Trace ID: UT-REC-CPP-02 +// Architecture: ARCH-FAC-001 +// Unit Design: UD-REC-01 +// Intent: RecurrentPipeline with Mamba, RWKV, and Hybrid state managers +// Preconditions: TRT + CUDA GPU available, mock engines +// Postconditions: Pipeline generates tokens with correct state management per backend type +// ============================================================================= + +// ============================================================================= +// Test suite: RecurrentPipeline (Mamba, RWKV, Hybrid) +// ============================================================================= +// +// Tests the RecurrentPipeline with mock engines and Qwen3.8-owned recurrent +// and hybrid state implementations. +// ============================================================================= + +#include "runtime/models/qwen3_8/hybrid_state.h" +#include "runtime/models/qwen3_8/kv_cache.h" +#include "runtime/models/qwen3_8/pipeline.h" +#include "runtime/models/qwen3_8/recurrent_state.h" +#include "trtmc/runtime/trt_module.h" +// pipeline_interface.h was removed; GenerateConfig is in trtmc/pipeline.h +// (already included transitively via recurrent_pipeline.h) + +#include "runtime/backend/trt_module_impl.h" +#include "runtime/core/trt_common.h" + +#include +#include +#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; + +class RecordingTokenizer final : public trtmc::ITokenizer { + public: + std::vector encode(const std::string& text) const override { + last_text = text; + return {1, 0}; + } + + 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 == "") + return 1; + if (token == "Paris") + return 2; + return 0; + } + + std::string token_for_id(int32_t id) const override { + if (id == 2) + return "Paris"; + return ""; + } + + mutable std::string last_text; +}; + +// Mock decoder: token_id[1] → logits[4] = constant [0.1, 0.2, 0.9, 0.3] +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); + + auto* inp = network->addInput("token_id", nvinfer1::DataType::kINT32, nvinfer1::Dims{1, {1}}); + float const_logits[4] = {0.1f, 0.2f, 0.9f, 0.3f}; + auto* cst = network->addConstant( + nvinfer1::Dims{1, {4}}, nvinfer1::Weights{nvinfer1::DataType::kFLOAT, const_logits, 4}); + cst->getOutput(0)->setName("logits"); + network->markOutput(*cst->getOutput(0)); + + // Use input so it's not optimized away + auto* id = network->addIdentity(*inp); + id->getOutput(0)->setName("_unused"); + + auto plan = trtmc::TrtUniquePtr( + builder->buildSerializedNetwork(*network, *config)); + if (!plan) + return nullptr; + auto rt = trtmc::TrtUniquePtr(nvinfer1::createInferRuntime(g_logger)); + return trtmc::TrtUniquePtr( + rt->deserializeCudaEngine(plan->data(), plan->size())); +} + +static void test_mamba_pipeline() { + auto engine = build_mock_decoder(); + if (!engine) { + std::cerr << "SKIP: can't build engine\n"; + return; + } + + cudaStream_t stream; + cudaStreamCreate(&stream); + + auto module = std::make_unique(engine.get(), + engine->createExecutionContext(), stream); + + // Mamba: 2 state specs, 1 layer + std::vector specs = { + {"conv_state", {12}}, + {"ssm_state", {32}}, + }; + auto rs = std::make_unique(1, specs, stream); + + trtmc::RecurrentGenConfig cfg; + cfg.vocab_size = 4; + cfg.id_eos = 2; // argmax=2=eos + + // Scoped so the pipeline destructor runs while `stream` is still valid. + { + trtmc::RecurrentPipeline pipeline(std::move(module), std::move(rs), cfg, stream, + "MambaPipeline"); + check(std::string(pipeline.pipeline_type()) == "MambaPipeline", "mamba name"); + + trtmc::GenerateConfig gen_cfg; + gen_cfg.max_new_tokens = 5; + auto result = pipeline.generate_ids({1}, gen_cfg); + + // argmax=2=eos → stops after 1 generated token + check(result.token_ids.size() == 2, "mamba: input + 1 generated"); + check(result.token_ids[1] == 2, "mamba: generated token = 2 (eos)"); + } + cudaStreamDestroy(stream); +} + +static void test_rwkv_pipeline() { + auto engine = build_mock_decoder(); + if (!engine) { + std::cerr << "SKIP: can't build engine\n"; + return; + } + + cudaStream_t stream; + cudaStreamCreate(&stream); + + auto module = std::make_unique(engine.get(), + engine->createExecutionContext(), stream); + + // RWKV: 5 state specs, 2 layers + std::vector specs = { + {"attn_state", {8}}, {"ff_state", {8}}, {"num_state", {8}}, + {"den_state", {8}}, {"max_state", {8}}, + }; + auto rs = std::make_unique(2, specs, stream); + + trtmc::RecurrentGenConfig cfg; + cfg.vocab_size = 4; + cfg.id_eos = 99; // never hit + + // Scoped so the pipeline destructor runs while `stream` is still valid. + { + trtmc::RecurrentPipeline pipeline(std::move(module), std::move(rs), cfg, stream, + "RwkvPipeline"); + check(std::string(pipeline.pipeline_type()) == "RwkvPipeline", "rwkv name"); + + trtmc::GenerateConfig gen_cfg; + gen_cfg.max_new_tokens = 3; + auto result = pipeline.generate_ids({0}, gen_cfg); + + check(result.token_ids.size() == 4, "rwkv: input + 3 generated"); + check(result.token_ids[1] == 2, "rwkv: all gen tokens = 2"); + check(result.token_ids[3] == 2, "rwkv: last gen = 2"); + } + cudaStreamDestroy(stream); +} + +static void test_hybrid_pipeline() { + auto engine = build_mock_decoder(); + if (!engine) { + std::cerr << "SKIP: can't build engine\n"; + return; + } + + cudaStream_t stream; + cudaStreamCreate(&stream); + + // Build mock engine with mask input too + auto builder = trtmc::TrtUniquePtr(nvinfer1::createInferBuilder(g_logger)); + auto network = trtmc::TrtUniquePtr(builder->createNetworkV2(0)); + auto bconfig = trtmc::TrtUniquePtr(builder->createBuilderConfig()); + bconfig->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, 1 << 20); + + auto* tok = network->addInput("token_id", nvinfer1::DataType::kINT32, nvinfer1::Dims{1, {1}}); + auto* pos = + network->addInput("position_id", nvinfer1::DataType::kINT32, nvinfer1::Dims{1, {1}}); + // Qwen38KvCache(1, /*max_length=*/4, ...) writes a decode mask of width + // max_length + 1 = 5. Declaring 4 here would silently drop the final column. + auto* mask = + network->addInput("attention_mask", nvinfer1::DataType::kFLOAT, nvinfer1::Dims{1, {5}}); + + float cl[4] = {0.1f, 0.2f, 0.9f, 0.3f}; + auto* c = network->addConstant(nvinfer1::Dims{1, {4}}, + nvinfer1::Weights{nvinfer1::DataType::kFLOAT, cl, 4}); + c->getOutput(0)->setName("logits"); + network->markOutput(*c->getOutput(0)); + + network->addIdentity(*tok)->getOutput(0)->setName("_t"); + network->addIdentity(*pos)->getOutput(0)->setName("_p"); + network->addIdentity(*mask)->getOutput(0)->setName("_m"); + + auto plan = trtmc::TrtUniquePtr( + builder->buildSerializedNetwork(*network, *bconfig)); + if (!plan) { + std::cerr << "SKIP: can't build engine\n"; + cudaStreamDestroy(stream); + return; + } + auto rt = trtmc::TrtUniquePtr(nvinfer1::createInferRuntime(g_logger)); + auto hybrid_engine = trtmc::TrtUniquePtr( + rt->deserializeCudaEngine(plan->data(), plan->size())); + if (!hybrid_engine) { + std::cerr << "SKIP: can't build engine\n"; + cudaStreamDestroy(stream); + return; + } + + auto module = std::make_unique( + hybrid_engine.get(), hybrid_engine->createExecutionContext(), stream); + auto kv = std::make_unique(1, 4, 2, stream); + std::vector specs = {{"ssm", {4}}}; + auto ssm = std::make_unique(1, specs, stream); + auto hybrid = std::make_unique(std::move(kv), std::move(ssm)); + + trtmc::RecurrentGenConfig cfg; + cfg.vocab_size = 4; + cfg.id_eos = 2; + cfg.has_position_input = true; + + // Scoped so the pipeline destructor runs while `stream` is still valid. + { + trtmc::RecurrentPipeline pipeline(std::move(module), std::move(hybrid), cfg, stream, + "HybridPipeline"); + check(std::string(pipeline.pipeline_type()) == "HybridPipeline", "hybrid name"); + + trtmc::GenerateConfig gen_cfg; + gen_cfg.max_new_tokens = 5; + auto result = pipeline.generate_ids({0}, gen_cfg); + + check(result.token_ids.size() == 2, "hybrid: input + eos"); + check(result.token_ids[1] == 2, "hybrid: eos generated"); + } + cudaStreamDestroy(stream); +} + +static void test_generate_applies_chat_template() { + auto engine = build_mock_decoder(); + if (!engine) { + std::cerr << "SKIP: can't build engine\n"; + return; + } + + cudaStream_t stream; + cudaStreamCreate(&stream); + + auto module = std::make_unique(engine.get(), + engine->createExecutionContext(), stream); + std::vector specs = { + {"conv_state", {12}}, + {"ssm_state", {32}}, + }; + auto rs = std::make_unique(1, specs, stream); + + trtmc::RecurrentGenConfig cfg; + cfg.vocab_size = 4; + cfg.id_bos = 1; + cfg.id_eos = 2; + cfg.chat_template_format = "chatml"; + + auto tokenizer = std::make_shared(); + // Scoped so the pipeline destructor runs while `stream` is still valid. + { + trtmc::RecurrentPipeline pipeline(std::move(module), std::move(rs), cfg, stream, + "MambaPipeline", tokenizer); + + trtmc::GenerateConfig gen_cfg; + gen_cfg.max_new_tokens = 1; + gen_cfg.use_chat_template = true; + gen_cfg.enable_thinking = false; + auto result = pipeline.generate("What is the capital of France?", gen_cfg); + + check(result.text == "Paris", "chat template: generated text decodes"); + check(tokenizer->last_text.find("<|im_start|>user\n") != std::string::npos, + "chat template: user prefix"); + check(tokenizer->last_text.find("What is the capital of France?") != std::string::npos, + "chat template: prompt retained"); + check(tokenizer->last_text.find("\n\n\n\n") != std::string::npos, + "chat template: no-thinking block is closed"); + } + cudaStreamDestroy(stream); +} + +static void test_argmax_recurrent() { + std::vector v = {-1.0f, 5.0f, 3.0f}; + check(trtmc::RecurrentPipeline::argmax(v) == 1, "argmax = 1"); +} + +int main() { + test_argmax_recurrent(); + test_mamba_pipeline(); + test_rwkv_pipeline(); + test_hybrid_pipeline(); + test_generate_applies_chat_template(); + if (failures > 0) + std::cerr << failures << " FAILED\n"; + return failures; +} diff --git a/tests/cpp/models/qwen3_8/test_qwen3_8_runtime_config_contract.cpp b/tests/cpp/models/qwen3_8/test_qwen3_8_runtime_config_contract.cpp new file mode 100644 index 000000000..c03558eef --- /dev/null +++ b/tests/cpp/models/qwen3_8/test_qwen3_8_runtime_config_contract.cpp @@ -0,0 +1,122 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +// CPU-only consumer contract for Qwen3.8's serialized runtime config. + +#include "trtmc/runtime/pipeline_plugin.h" + +#include +#include +#include +#include + +int main() { + // Transcribed from the config.json section of a bundle built from + // Qwen/Qwen3.8-27B at 1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0. Qwen3.8 + // keeps every decoder dimension under text_config, but the runtime reads + // the bundle with a top-level nlohmann lookup, so the duplicated top-level + // fields are the actual contract. Without them compute_kv_dim() returns 0 + // and the KV cache allocates zero-sized tensors. + const std::string config = R"({ + "vocab_size": 248320, + "hidden_size": 5120, + "num_hidden_layers": 64, + "num_attention_heads": 24, + "num_key_value_heads": 4, + "head_dim": 256, + "intermediate_size": 17408, + "max_position_embeddings": 262144, + "rms_norm_eps": 1e-06, + "bos_token_id": 248044, + "model_type": "qwen3_5", + "text_config": { + "model_type": "qwen3_5_text", + "vocab_size": 248320, + "hidden_size": 5120, + "intermediate_size": 17408, + "num_hidden_layers": 64, + "num_attention_heads": 24, + "num_key_value_heads": 4, + "head_dim": 256, + "bos_token_id": 248044, + "eos_token_id": 248044, + "output_gate_type": "swish", + "max_position_embeddings": 262144, + "linear_conv_kernel_dim": 4, + "linear_key_head_dim": 128, + "linear_num_key_heads": 16, + "linear_num_value_heads": 48, + "linear_value_head_dim": 128, + "rope_parameters": { + "mrope_interleaved": true, + "mrope_section": [11, 11, 10], + "rope_type": "default", + "rope_theta": 10000000, + "partial_rotary_factor": 0.25 + } + }, + "vision_config": { + "model_type": "qwen3_5", + "depth": 27, + "hidden_size": 1152, + "intermediate_size": 4304, + "num_heads": 16, + "out_hidden_size": 5120, + "patch_size": 16, + "spatial_merge_size": 2, + "temporal_patch_size": 2, + "deepstack_visual_indexes": [] + }, + "num_mamba_layers": 48, + "num_attention_layers": 16, + "d_inner": 6144, + "mamba_d_state": 128, + "mamba_d_conv": 4, + "mamba_nheads": 48, + "mamba_head_dim": 128, + "conv_dim": 10240, + "eos_token_id": [248046, 248044], + "runtime_strategy": "qwen3_8_hybrid_mamba_attention", + "precision": "fp16", + "tokenizer_add_special_tokens": 0 + })"; + + const auto parsed = trtmc::parse_base_config(config, 256); + bool ok = true; + const auto check = [&](bool condition, const char* name) { + if (!condition) { + std::cerr << "FAIL: " << name << '\n'; + ok = false; + } + }; + + check(parsed.runtime_strategy == "qwen3_8_hybrid_mamba_attention", "runtime strategy"); + check(parsed.precision == "fp16", "precision"); + check(parsed.vocab_size == 248320, "vocabulary size"); + check(parsed.hidden_size == 5120, "hidden size"); + check(parsed.num_layers == 64, "layer count"); + check(parsed.num_heads == 24, "attention head count"); + check(parsed.num_kv_heads == 4, "KV head count"); + check(parsed.head_dim == 256, "head dimension"); + check(parsed.attention_size == 6144, "attention width"); + check(parsed.num_kv_heads * parsed.head_dim == 1024, "KV cache width"); + check(parsed.max_cache_length == 256, "cache length override"); + check(parsed.id_bos == 248044, "BOS token"); + + // Qwen3.8 diverges from Qwen3.5 here. text_config carries a single + // eos_token_id (248044), but the checkpoint actually terminates on 248046, + // which only appears in generation_config.json. The builder serializes that + // full list, and the family deliberately does not republish the + // text_config value as a bundle override, because overrides are merged last + // and would collapse the list to 248044 alone -- leaving 248046 unmatched + // and generation running to max_new_tokens. + check(parsed.id_eos == 248046, "primary EOS token from generation config"); + check(parsed.id_eos_ids == std::vector({248046, 248044}), + "full EOS token list survives override merge"); + + check(parsed.tokenizer_add_special_tokens_present, "tokenizer flag presence"); + check(!parsed.tokenizer_add_special_tokens, "tokenizer flag value"); + return ok ? 0 : 1; +} diff --git a/tests/e2e/models/qwen3_8/MODEL.toml b/tests/e2e/models/qwen3_8/MODEL.toml new file mode 100644 index 000000000..98d13660e --- /dev/null +++ b/tests/e2e/models/qwen3_8/MODEL.toml @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +id = "qwen3_8" +plugin = "qwen3_8" +test_manifests = [ + "manifests/qwen38-27b.json", +] + +[e2e_defaults.text_generation_causal] +reference_backend = "hf_transformers" +oracle_level = "L1_external_reference" +stages = [ + { name = "full_generation", required = true }, +] diff --git a/tests/e2e/models/qwen3_8/e2e_plugins/__init__.py b/tests/e2e/models/qwen3_8/e2e_plugins/__init__.py new file mode 100644 index 000000000..1f0515e3f --- /dev/null +++ b/tests/e2e/models/qwen3_8/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/qwen3_8/e2e_plugins/comparator.py b/tests/e2e/models/qwen3_8/e2e_plugins/comparator.py new file mode 100644 index 000000000..dec0325b0 --- /dev/null +++ b/tests/e2e/models/qwen3_8/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 + +"""qwen3_8 model-owned E2E comparator plugins.""" + +from __future__ import annotations + +from .comparators.text import TextComparator + + +class Qwen38TextGenerationCausalComparator(TextComparator): + """qwen3_8 local comparator for text_generation_causal.""" + +comparator = Qwen38TextGenerationCausalComparator() diff --git a/tests/e2e/models/qwen3_8/e2e_plugins/comparators/__init__.py b/tests/e2e/models/qwen3_8/e2e_plugins/comparators/__init__.py new file mode 100644 index 000000000..9c9857522 --- /dev/null +++ b/tests/e2e/models/qwen3_8/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/qwen3_8/e2e_plugins/comparators/_helpers.py b/tests/e2e/models/qwen3_8/e2e_plugins/comparators/_helpers.py new file mode 100644 index 000000000..20d4d897b --- /dev/null +++ b/tests/e2e/models/qwen3_8/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/qwen3_8/e2e_plugins/comparators/text.py b/tests/e2e/models/qwen3_8/e2e_plugins/comparators/text.py new file mode 100644 index 000000000..016ea55c4 --- /dev/null +++ b/tests/e2e/models/qwen3_8/e2e_plugins/comparators/text.py @@ -0,0 +1,526 @@ +# 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}", + ) + + # Step counts must agree. Silently truncating to the shorter side + # discards every unmatched decode step, so an early EOS or a dropped + # step would be scored only on the shared prefix and could pass. + if trt_logits.shape[0] != ref_logits.shape[0]: + return CompareResult( + stage_name=stage.name, + status=StageStatus.FAILED.value, + metrics=metrics, + message=( + "Generation length mismatch: " + f"TRT produced {trt_logits.shape[0]} decode steps, " + f"reference produced {ref_logits.shape[0]}" + ), + ) + + n_steps = trt_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) + # Only the documented direction is forgiven: TRT stopping early on + # EOS while the reference keeps emitting. A reference shorter than + # the TRT output is not that case and must not be excused. + if (token_agreement_rate >= ta_thresh + and len(trt_text_for_ned) <= len(ref_text_for_ned)): + short, long = trt_text_for_ned, ref_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/qwen3_8/e2e_plugins/contract.py b/tests/e2e/models/qwen3_8/e2e_plugins/contract.py new file mode 100644 index 000000000..68086d9fa --- /dev/null +++ b/tests/e2e/models/qwen3_8/e2e_plugins/contract.py @@ -0,0 +1,191 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Qwen3.8-owned multimodal chat contract plugin.""" + +from __future__ import annotations + +from tests.e2e_harness.contracts import MetricResult +# Model-owned contract helpers. Keep behavior here so contract semantics do not +# drift across model families through shared harness code. +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] + import re + 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 Qwen38MultimodalChatPlugin: + reference_families = ["multimodal_chat_qwen38"] + user_contract = "chat_response" + + def configure_reference(self, case): + return { + "use_chat_template": True, + "use_processor": True, + "enable_thinking": False, + } + + def verify(self, trt_output, ref_output, case, threshold): + prompt = case.inputs.get("prompt", "") + trt_answer = normalize_text(extract_answer(trt_output, prompt)) + ref_answer = normalize_text(extract_answer(ref_output, prompt)) + + if not trt_answer: + return make_fail("full_generation", {}, message="TRT produced empty response") + + exact_match = trt_answer == ref_answer + ned = levenshtein_ned(trt_answer, ref_answer) + ned_threshold = threshold.metrics.get("contract_ned_threshold", 0.15) + metrics = { + "exact_match": MetricResult( + value=1.0 if exact_match else 0.0, + threshold=1.0, + operator="==", + passed=exact_match, + ), + "ned": MetricResult( + value=ned, + threshold=ned_threshold, + operator="<=", + passed=ned <= ned_threshold, + ), + } + + rule = "exact_match OR ned <= threshold" + if exact_match or ned <= ned_threshold: + return make_pass("full_generation", metrics, rule) + return make_fail( + "full_generation", + metrics, + rule, + f"Qwen3.8 multimodal chat response diverged: NED={ned:.3f}", + ) + +plugin = Qwen38MultimodalChatPlugin() diff --git a/tests/e2e/models/qwen3_8/e2e_plugins/contracts.py b/tests/e2e/models/qwen3_8/e2e_plugins/contracts.py new file mode 100644 index 000000000..d6f9281d2 --- /dev/null +++ b/tests/e2e/models/qwen3_8/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/qwen3_8/e2e_plugins/reference.py b/tests/e2e/models/qwen3_8/e2e_plugins/reference.py new file mode 100644 index 000000000..dba07bd93 --- /dev/null +++ b/tests/e2e/models/qwen3_8/e2e_plugins/reference.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""qwen3_8 model-owned E2E reference plugins.""" + +from __future__ import annotations + +from .references.hf_transformers import HfTransformersReference + + +class Qwen38HfTransformersReference(HfTransformersReference): + """qwen3_8 local reference for hf_transformers.""" + +reference = Qwen38HfTransformersReference() diff --git a/tests/e2e/models/qwen3_8/e2e_plugins/references/__init__.py b/tests/e2e/models/qwen3_8/e2e_plugins/references/__init__.py new file mode 100644 index 000000000..6487ae800 --- /dev/null +++ b/tests/e2e/models/qwen3_8/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/qwen3_8/e2e_plugins/references/hf_transformers.py b/tests/e2e/models/qwen3_8/e2e_plugins/references/hf_transformers.py new file mode 100644 index 000000000..59376ddc1 --- /dev/null +++ b/tests/e2e/models/qwen3_8/e2e_plugins/references/hf_transformers.py @@ -0,0 +1,1033 @@ +# 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, revision: str = "") -> str: + """Prefer a locally cached HF snapshot to avoid Hub API rate limits. + + A manifest that pins ``hf_revision`` must resolve that exact commit, not + whatever ``main`` currently points at; the reference runs with + ``trust_remote_code``, so the revision decides which code executes. + """ + if not hf_id: + return hf_id + p = Path(hf_id) + if p.exists(): + return hf_id + + try: + from huggingface_hub import snapshot_download + + kwargs = {"local_files_only": True} + if revision: + kwargs["revision"] = revision + return snapshot_download(hf_id, **kwargs) + 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 + # The child imports this module by its full dotted path, so the repository + # root has to be importable even when the parent process runs from another + # directory. Existing entries are preserved. + project_dir = str(PROJECT_DIR) + existing = env.get("PYTHONPATH", "") + parts = [p for p in existing.split(os.pathsep) if p] + if project_dir not in parts: + parts.insert(0, project_dir) + env["PYTHONPATH"] = os.pathsep.join(parts) + 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, case.hf_revision) + # A resolved snapshot path is already the pinned commit; a bare hf_id is + # not, so the revision must travel with the loader call. + revision_kwargs = ( + {"revision": case.hf_revision} + if case.hf_revision and not Path(model_ref).exists() else {}) + 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} + revision_kwargs = {revision_kwargs!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, **revision_kwargs) + 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}, + **revision_kwargs, + }} + # Detect encoder-decoder models by checking config + from transformers import AutoConfig + _cfg = AutoConfig.from_pretrained( + model_ref, trust_remote_code=trust_remote_code, **revision_kwargs) + 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, case.hf_revision) + # A resolved snapshot path is already the pinned commit; a bare hf_id is + # not, so the revision must travel with the loader call. + revision_kwargs = ( + {"revision": case.hf_revision} + if case.hf_revision and not Path(model_ref).exists() else {}) + 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} + revision_kwargs = {revision_kwargs!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, case.hf_revision) + # A resolved snapshot path is already the pinned commit; a bare hf_id is + # not, so the revision must travel with the loader call. + revision_kwargs = ( + {"revision": case.hf_revision} + if case.hf_revision and not Path(model_ref).exists() else {}) + 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} + revision_kwargs = {revision_kwargs!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, case.hf_revision) + # A resolved snapshot path is already the pinned commit; a bare hf_id is + # not, so the revision must travel with the loader call. + revision_kwargs = ( + {"revision": case.hf_revision} + if case.hf_revision and not Path(model_ref).exists() else {}) + 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} + revision_kwargs = {revision_kwargs!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, case.hf_revision) + # A resolved snapshot path is already the pinned commit; a bare hf_id is + # not, so the revision must travel with the loader call. + revision_kwargs = ( + {"revision": case.hf_revision} + if case.hf_revision and not Path(model_ref).exists() else {}) + 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} + revision_kwargs = {revision_kwargs!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/qwen3_8/e2e_plugins/runner.py b/tests/e2e/models/qwen3_8/e2e_plugins/runner.py new file mode 100644 index 000000000..6143e8c55 --- /dev/null +++ b/tests/e2e/models/qwen3_8/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 + +"""qwen3_8 model-owned E2E runner plugins.""" + +from __future__ import annotations + +from .runners.text_generation import TextGenerationCausalRunner + + +class Qwen38TextGenerationCausalRunner(TextGenerationCausalRunner): + """qwen3_8 local runner for text_generation_causal.""" + +runner = Qwen38TextGenerationCausalRunner() diff --git a/tests/e2e/models/qwen3_8/e2e_plugins/runners/__init__.py b/tests/e2e/models/qwen3_8/e2e_plugins/runners/__init__.py new file mode 100644 index 000000000..986743e50 --- /dev/null +++ b/tests/e2e/models/qwen3_8/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/qwen3_8/e2e_plugins/runners/text_generation.py b/tests/e2e/models/qwen3_8/e2e_plugins/runners/text_generation.py new file mode 100644 index 000000000..aebf84ddb --- /dev/null +++ b/tests/e2e/models/qwen3_8/e2e_plugins/runners/text_generation.py @@ -0,0 +1,853 @@ +# 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 decoder_kv_cache, decoder_moe, mamba_ssm_recurrent, rwkv_recurrent, and +hybrid_mamba_attention runtime strategies, all of which map 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 +from ..runtime_config import runtime_config_set_tokens + +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 + try: + sample = json.loads(line) + except json.JSONDecodeError: + # A truncated or non-JSON first line means the runner produced + # no usable sample. Returning {} lets the caller fall back to + # stdout and report a failed case with captured stderr, instead + # of aborting the whole stage on an unhandled exception. + return {} + if not isinstance(sample, dict): + return {} + token_ids = sample.get("token_ids") + if isinstance(token_ids, list): + sample["token_ids"] = [int(token) for token in token_ids] + return sample + return {} + + +def _ensure_distributed_runtime_env( + case: E2ECase, + ctx: RunContext, + env: dict[str, str], + rendezvous_suffix: str = "", +) -> None: + """Populate shared env values needed by all distributed ranks.""" + if not _distributed_runtime_config(case): + return + if env.get("TRTMC_NCCL_RENDEZVOUS"): + return + + safe_name = _safe_artifact_name(case.name) + root = Path(_case_artifact_dir(ctx.artifacts_dir, case.name)) if ctx.artifacts_dir else \ + Path(tempfile.gettempdir()) + path = root / f"{safe_name}{rendezvous_suffix}.nccl_rendezvous.bin" + path.parent.mkdir(parents=True, exist_ok=True) + try: + path.unlink() + except FileNotFoundError: + pass + env["TRTMC_NCCL_RENDEZVOUS"] = str(path) + + +def _wrap_distributed_command( + cmd: list[str], case: E2ECase | None, env: dict[str, str] +) -> list[str]: + config = _distributed_runtime_config(case) + if not config: + return cmd + + launcher = str(config.get("launcher", "mpirun") or "mpirun") + world_size = int(config.get("world_size", config.get("tp_size", 2)) or 2) + launcher_args = config.get("launcher_args") + if isinstance(launcher_args, list): + prefix = [launcher] + [str(arg) for arg in launcher_args] + else: + prefix = [launcher, "--tag-output", "-np", str(world_size)] + + export_env = config.get("export_env", ["LD_LIBRARY_PATH", "CUDA_VISIBLE_DEVICES"]) + if isinstance(export_env, list) and Path(launcher).name == "mpirun": + export_names = [str(name) for name in export_env] + for name in ("TRTMC_NCCL_RENDEZVOUS", "TRTMC_EMBEDDING_STDOUT"): + if name in env and name not in export_names: + export_names.append(name) + for name in export_names: + if name in env: + prefix.extend(["-x", name]) + + return prefix + cmd + + +def _visible_gpu_indices(env: dict[str, str]) -> list[str]: + raw = env.get("CUDA_VISIBLE_DEVICES", "") + if not raw or raw.lower() in {"all", "none", "void"}: + return [] + indices: list[str] = [] + for part in raw.split(","): + token = part.strip() + if token.isdigit(): + indices.append(token) + return indices + + +class _GpuMemorySampler: + def __init__(self, artifacts_dir: str | None, case_name: str, env: dict[str, str], + interval_ms: int) -> None: + root = Path(_case_artifact_dir(artifacts_dir, case_name)) if artifacts_dir else \ + Path(tempfile.gettempdir()) + root.mkdir(parents=True, exist_ok=True) + self.path = root / "gpu_memory_samples.csv" + self.env = env + self.interval_ms = max(50, interval_ms) + self.visible_indices = _visible_gpu_indices(env) + self.proc: subprocess.Popen | None = None + self.handle = None + self.error = "" + + def start(self) -> None: + if shutil.which("nvidia-smi") is None: + self.error = "nvidia-smi not found" + return + self.handle = self.path.open("w", encoding="utf-8") + cmd = [ + "nvidia-smi", + "--query-gpu=index,memory.used", + "--format=csv,noheader,nounits", + f"--loop-ms={self.interval_ms}", + ] + try: + self.proc = subprocess.Popen( + cmd, + stdout=self.handle, + stderr=subprocess.DEVNULL, + text=True, + env=self.env, + ) + except Exception as exc: + self.error = str(exc) + self.handle.close() + self.handle = None + + def stop(self) -> dict: + if self.proc is not None: + self.proc.terminate() + try: + self.proc.wait(timeout=2) + except subprocess.TimeoutExpired: + self.proc.kill() + self.proc.wait(timeout=2) + if self.handle is not None: + self.handle.close() + self.handle = None + return self._summary() + + def _summary(self) -> dict: + meta = { + "sample_file": str(self.path), + "sample_interval_ms": self.interval_ms, + "visible_device_indices": self.visible_indices, + } + if self.error: + meta["error"] = self.error + return meta + peaks: dict[str, int] = {} + sample_count = 0 + if not self.path.is_file(): + meta["error"] = "sample file was not created" + return meta + with self.path.open("r", encoding="utf-8") as f: + for line in f: + parts = [p.strip() for p in line.split(",")] + if len(parts) < 2 or not parts[0].isdigit(): + continue + if self.visible_indices and parts[0] not in self.visible_indices: + continue + try: + used_mb = int(float(parts[1])) + except ValueError: + continue + peaks[parts[0]] = max(peaks.get(parts[0], 0), used_mb) + sample_count += 1 + meta["sample_count"] = sample_count + meta["peak_memory_mb_by_gpu"] = peaks + if peaks: + meta["peak_memory_mb"] = max(peaks.values()) + meta["peak_memory_mb_visible_sum"] = sum(peaks.values()) + return meta + + +def _maybe_start_gpu_memory_sampler( + distributed_runtime: dict, ctx: RunContext, case: E2ECase | None, env: dict[str, str] +) -> _GpuMemorySampler | None: + if case is None or not distributed_runtime.get("capture_gpu_memory"): + return None + interval_ms = int(distributed_runtime.get("gpu_memory_sample_interval_ms", 200) or 200) + sampler = _GpuMemorySampler(ctx.artifacts_dir, case.name, env, interval_ms) + sampler.start() + return sampler + + +def _extract_trtmc_timing(stderr: str) -> dict[str, float]: + match = _TRTMC_TIMING_RE.search(stderr or "") + if match is None: + return {} + try: + prefill_ms = float(match.group("prefill_ms")) + decode_ms = float(match.group("decode_ms")) + total_ms = float(match.group("total_ms")) + except ValueError: + return {} + return { + "trt_engine_prefill_s": prefill_ms / 1000.0, + "trt_engine_decode_s": decode_ms / 1000.0, + "trt_engine_s": total_ms / 1000.0, + } + + +def _extract_trtmc_load_timing(stderr: str) -> dict[str, float]: + total_ms = 0.0 + found = False + for match in _TRTMC_LOAD_TIMING_RE.finditer(stderr or ""): + try: + total_ms += float(match.group("load_deserialize_ms")) + found = True + except ValueError: + continue + return {"trt_load_deserialize_s": total_ms / 1000.0} if found else {} + + +def _detect_trt_runtime_error(stderr: str) -> str: + match = _TRT_RUNTIME_ERROR_RE.search(stderr or "") + return match.group(0).strip() if match else "" + + +def _distributed_debug_logits_required(case: E2ECase) -> bool: + distributed_runtime = _distributed_runtime_config(case) + return bool(distributed_runtime and distributed_runtime.get("debug_logits", True)) + + +def _format_debug_runner_error(case: E2ECase, phase: str, meta: dict) -> str: + detail = meta.get("error") + if not detail and meta.get("returncode") not in (None, 0): + detail = f"returncode={meta['returncode']}" + if not detail: + detail = "logits were not produced" + log_path = meta.get("stderr_log") + if log_path: + detail = f"{detail}; stderr_log={log_path}" + return ( + f"Distributed debug logits requested for {case.name} phase={phase}, " + f"but {detail}" + ) + + +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 = case.inputs.get("prompt", "The capital of France is") + max_new_tokens = case.inputs.get("max_new_tokens", 30) + 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 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, + }, + ) + + # 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 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}, + ) + + # ------------------------------------------------------------------ + # 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 = case.inputs.get("prompt", "The capital of France is") + + 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 = case.inputs.get("prompt", "The capital of France is") + 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") + for token in runtime_config_set_tokens(case): + cmd.extend(["--set", token]) + + 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", case.name) + 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" + ) + # Fixed path per case and phase: drop any file from an earlier run so a + # subprocess that exits 0 without writing cannot pass on stale logits. + Path(logits_path).unlink(missing_ok=True) + + 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.qwen3_8.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/qwen3_8/e2e_plugins/runtime_config.py b/tests/e2e/models/qwen3_8/e2e_plugins/runtime_config.py new file mode 100644 index 000000000..f9e763100 --- /dev/null +++ b/tests/e2e/models/qwen3_8/e2e_plugins/runtime_config.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Helpers for manifest-driven runtime config overrides.""" + +from __future__ import annotations + +from collections.abc import Iterator +from typing import Any + +from .contracts import E2ECase + + +def _runtime_config(case: E2ECase) -> dict[str, Any]: + config = case.metadata.get("runtime_config") + if isinstance(config, dict): + return config + config = case.inputs.get("runtime_config") + if isinstance(config, dict): + return config + return {} + + +def _flatten(prefix: str, value: Any) -> Iterator[tuple[str, Any]]: + if isinstance(value, dict): + for key, nested in value.items(): + name = f"{prefix}.{key}" if prefix else str(key) + yield from _flatten(name, nested) + elif prefix: + yield prefix, value + + +def _format_value(value: Any) -> str: + if isinstance(value, bool): + return "true" if value else "false" + return str(value) + + +def runtime_config_set_tokens(case: E2ECase) -> list[str]: + """Return CLI --set tokens from a manifest runtime_config mapping.""" + return [f"{name}={_format_value(value)}" for name, value in _flatten("", _runtime_config(case))] + + +def runtime_config_get(case: E2ECase, dotted_name: str, default: Any = None) -> Any: + value: Any = _runtime_config(case) + for part in dotted_name.split("."): + if not isinstance(value, dict) or part not in value: + return default + value = value[part] + return value diff --git a/tests/e2e/models/qwen3_8/manifests/qwen38-27b.json b/tests/e2e/models/qwen3_8/manifests/qwen38-27b.json new file mode 100644 index 000000000..7b89330fd --- /dev/null +++ b/tests/e2e/models/qwen3_8/manifests/qwen38-27b.json @@ -0,0 +1,30 @@ +{ + "name": "qwen38-27b", + "hf_id": "Qwen/Qwen3.8-27B", + "hf_revision": "1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0", + "bundle": "qwen38-27b.bundle", + "family": "qwen3_8", + "runtime_strategy": "qwen3_8_hybrid_mamba_attention", + "task_strategy": "text_generation_causal", + "e2e_parallel_resource": "exclusive_gpu", + "max_cache_length": 256, + "trust_remote_code": true, + "precision": "fp16", + "testcases": [ + { + "name": "qwen38-27b", + "trace_id": "IT-E2E-Q38-01", + "reference_family": "multimodal_chat_qwen38", + "user_contract": "chat_response", + "prompt": "What is the capital of France? Answer in one word.", + "max_new_tokens": 10, + "reference_precision": "fp32", + "runtime_config": { + "runtime": { + "disable_cuda_graph": true + } + }, + "notes": "Qwen3.8-27B decodes in FP16 against an FP32 Hugging Face reference. The decoder is the 64-layer hybrid stack (48 gated DeltaNet + 16 full-attention); the checkpoint's vision tower and MTP head are not part of the graph." + } + ] +} diff --git a/tests/e2e/models/qwen3_8/runner.py b/tests/e2e/models/qwen3_8/runner.py new file mode 100644 index 000000000..455f96949 --- /dev/null +++ b/tests/e2e/models/qwen3_8/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 qwen3_8 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/qwen3_8/test_qwen3_8_debug_runner.py b/tests/e2e/models/qwen3_8/test_qwen3_8_debug_runner.py new file mode 100644 index 000000000..6de05786c --- /dev/null +++ b/tests/e2e/models/qwen3_8/test_qwen3_8_debug_runner.py @@ -0,0 +1,79 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Qwen3.8-owned debug runner dispatch tests.""" + +from __future__ import annotations + +import json +import struct +from unittest.mock import patch + + +def _make_bundle_bytes( + header: dict, + engine_plan: bytes = b"FAKE_ENGINE_PLAN", + extra_sections: dict[str, bytes] | None = None, +) -> bytes: + magic = b"BUNDLE\x01\x00" + sections: dict[str, dict] = {} + body = b"" + + sections["engine_plan"] = {"offset": len(body), "size": len(engine_plan)} + body += engine_plan + + if extra_sections: + for name, data in extra_sections.items(): + sections[name] = {"offset": len(body), "size": len(data)} + body += data + + header["sections"] = sections + header_json = json.dumps(header).encode("utf-8") + return magic + struct.pack(" None: + _runner.run_model_e2e(case_name, request) diff --git a/tests/e2e/models/qwen3_8/test_qwen3_8_family_plugin.py b/tests/e2e/models/qwen3_8/test_qwen3_8_family_plugin.py new file mode 100644 index 000000000..07ba648b9 --- /dev/null +++ b/tests/e2e/models/qwen3_8/test_qwen3_8_family_plugin.py @@ -0,0 +1,534 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Branch-focused tests for the Qwen3.8 family plugin. + +Trace: ARCH-FAM-001, UD-FAM-QWEN3-8 +Intent: Validate Qwen3.8 DeltaNet/attention layer routing normalization and weight loading branches +Preconditions: Mixed DeltaNet/attention layer types and synthetic tensors with partial keys are provided +Postconditions: Layer type aliases are normalized correctly and branch-specific weights load with fallback behavior +""" + +from __future__ import annotations + +import json +from unittest.mock import patch + +import numpy as np +import pytest + +trt = pytest.importorskip( + "tensorrt", reason="TensorRT is required for family builder tests" +) + + +try: + from tensorrt_model_connect.config import ModelConfig + import tensorrt_model_connect.families.qwen3_8 as qwen3_8 +except (ImportError, ModuleNotFoundError): + pytest.skip("tensorrt_model_connect requires tensorrt", allow_module_level=True) + + +def _seq(*shape: int, start: int = 0) -> np.ndarray: + size = int(np.prod(shape)) + return np.arange(start, start + size, dtype=np.float32).reshape(shape) + + +def _patch_tensor_io(monkeypatch: pytest.MonkeyPatch, + tensor_map: dict[str, np.ndarray]) -> None: + monkeypatch.setattr(qwen3_8, "_open_safetensors", lambda _: ["reader"]) + monkeypatch.setattr( + qwen3_8, "_has_tensor", lambda _readers, name: name in tensor_map) + + def _load(_readers, name: str): + if name not in tensor_map: + raise KeyError(name) + return tensor_map[name] + + monkeypatch.setattr(qwen3_8, "_load_tensor", _load) + + +def test_parse_layer_types_normalizes_aliases(): + """Intent: validate routing normalization for mixed aliases. + Preconditions: layer type strings include canonical, alias, and unknown values. + Postconditions: aliases map to deltanet/attention and unknowns are lower-cased. + """ + parsed = qwen3_8._parse_layer_types( + ["linear", "FULL", "linear_attention", "full_attention", "Custom"]) + assert parsed == ["deltanet", "attention", "deltanet", "attention", "custom"] + + +def test_fp16_runtime_inputs_preserve_fp32_recurrent_state(): + """FP16 storage must not quantize the persistent DeltaNet state.""" + + class _Tensor: + def __init__(self, name: str, dtype): + self.name = name + self.dtype = dtype + + class _Layer: + def __init__(self, output: _Tensor): + self.output = output + + def get_output(self, index: int) -> _Tensor: + assert index == 0 + return self.output + + class _Network: + def __init__(self): + self.cast_inputs: list[_Tensor] = [] + + def add_cast(self, tensor: _Tensor, dtype): + self.cast_inputs.append(tensor) + return _Layer(_Tensor(f"{tensor.name}_cast", dtype)) + + network = _Network() + attention_mask = _Tensor("attention_mask", trt.float32) + conv_state = _Tensor("conv_state", trt.float32) + ssm_state = _Tensor("ssm_state", trt.float32) + cache_k = _Tensor("cache_k", trt.float16) + cache_v = _Tensor("cache_v", trt.float16) + + ( + prepared_mask, + prepared_conv, + prepared_ssm, + prepared_cache_k, + prepared_cache_v, + ) = qwen3_8._prepare_runtime_inputs( + network, + trt.float16, + attention_mask, + [conv_state], + [ssm_state], + [cache_k], + [cache_v], + ) + + assert prepared_mask.dtype == trt.float16 + assert prepared_conv[0].dtype == trt.float16 + assert prepared_cache_k[0].dtype == trt.float16 + assert prepared_cache_v[0].dtype == trt.float16 + assert prepared_ssm == [ssm_state] + assert prepared_ssm[0].dtype == trt.float32 + assert ssm_state not in network.cast_inputs + + +def test_load_weights_mixed_branches_and_fallbacks(monkeypatch: pytest.MonkeyPatch): + """Intent: execute DeltaNet + attention branches and optional-key fallbacks. + Preconditions: one layer is DeltaNet, one layer is attention, with partial tensors missing. + Postconditions: normalized weights/metadata are emitted with correct fallback behavior. + """ + raw = { + "text_config": { + "layer_types": ["linear", "FULL"], + "linear_num_value_heads": 2, + "linear_num_key_heads": 1, + "linear_value_head_dim": 2, + "linear_conv_kernel_dim": 3, + "rope_parameters": { + "partial_rotary_factor": 0.5, + "rope_theta": 321.0, + }, + } + } + cfg = ModelConfig( + model_type="qwen3_8", + vocab_size=5, + hidden_size=8, + intermediate_size=6, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=1, + rope_theta=9999.0, + raw=raw, + ) + + tensors: dict[str, np.ndarray] = { + # Embedding only under fallback key. + "model.embed_tokens.weight": _seq(5, 8, start=0), + # Layer 0 DeltaNet path. + "model.language_model.layers.0.input_layernorm.weight": _seq(8, start=100), + "model.language_model.layers.0.linear_attn.in_proj_qkv.weight": _seq( + 8, 8, start=200 + ), + "model.language_model.layers.0.linear_attn.in_proj_z.weight": _seq( + 4, 8, start=400 + ), + "model.language_model.layers.0.linear_attn.in_proj_a.weight": _seq( + 2, 8, start=500 + ), + "model.language_model.layers.0.linear_attn.in_proj_b.weight": _seq( + 2, 8, start=600 + ), + "model.language_model.layers.0.linear_attn.A_log": np.log( + np.array([1.0, 2.0], dtype=np.float32) + ), + "model.language_model.layers.0.linear_attn.dt_bias": np.array( + [0.25, -0.5], dtype=np.float32 + ), + "model.language_model.layers.0.linear_attn.conv1d.weight": _seq( + 8, 1, 3, start=700 + ), + # conv1d.bias intentionally omitted to exercise zero fallback. + "model.language_model.layers.0.linear_attn.norm.weight": np.array( + [0.5, 1.5], dtype=np.float32 + ), + "model.language_model.layers.0.linear_attn.out_proj.weight": _seq( + 8, 4, start=800 + ), + "model.language_model.layers.0.mlp.gate_proj.weight": _seq(6, 8, start=900), + "model.language_model.layers.0.mlp.up_proj.weight": _seq(6, 8, start=1000), + "model.language_model.layers.0.mlp.down_proj.weight": _seq(8, 6, start=1100), + # Layer 1 attention path. + # input_layernorm intentionally omitted to exercise ones fallback. + "model.language_model.layers.1.post_attention_layernorm.weight": _seq( + 8, start=1200 + ), + "model.language_model.layers.1.self_attn.q_proj.weight": _seq( + 16, 8, start=1300 + ), + "model.language_model.layers.1.self_attn.k_proj.weight": _seq(4, 8, start=1500), + "model.language_model.layers.1.self_attn.v_proj.weight": _seq(4, 8, start=1600), + "model.language_model.layers.1.self_attn.o_proj.weight": _seq(8, 8, start=1700), + "model.language_model.layers.1.self_attn.q_norm.weight": np.array( + [0.1, 0.2, 0.3, 0.4], dtype=np.float32 + ), + # k_norm intentionally omitted. + # Layer 1 MLP intentionally omitted (gate check should skip all MLP loads). + } + _patch_tensor_io(monkeypatch, tensors) + + weights = qwen3_8.plugin.load_weights("/unused", cfg) + + np.testing.assert_allclose(weights["embedding"], tensors["model.embed_tokens.weight"]) + + np.testing.assert_allclose( + weights["layer.0.input_norm"], + 1.0 + tensors["model.language_model.layers.0.input_layernorm.weight"], + ) + np.testing.assert_allclose(weights["layer.1.input_norm"], np.ones(8, dtype=np.float32)) + + np.testing.assert_allclose( + weights["layer.0.post_attn_norm"], np.ones(8, dtype=np.float32)) + np.testing.assert_allclose( + weights["layer.1.post_attn_norm"], + 1.0 + tensors["model.language_model.layers.1.post_attention_layernorm.weight"], + ) + + np.testing.assert_allclose( + weights["layer.0.conv1d_bias"], np.zeros(8, dtype=np.float32)) + np.testing.assert_allclose( + weights["layer.0.deltanet_norm"], + np.array([0.5, 1.5, 0.5, 1.5], dtype=np.float32), + ) + np.testing.assert_allclose( + weights["layer.0.A"], np.array([-1.0, -2.0], dtype=np.float32)) + + q_raw = tensors["model.language_model.layers.1.self_attn.q_proj.weight"] + q_reshaped = q_raw.reshape(2, 8, 8) + expected_q = q_reshaped[:, :4, :].reshape(8, 8).T.astype(np.float32) + expected_gate = q_reshaped[:, 4:, :].reshape(8, 8).T.astype(np.float32) + np.testing.assert_allclose(weights["layer.1.w_q"], expected_q) + np.testing.assert_allclose(weights["layer.1.w_gate_attn"], expected_gate) + np.testing.assert_allclose( + weights["layer.1.q_norm"], + np.tile( + 1.0 + tensors["model.language_model.layers.1.self_attn.q_norm.weight"], + 2, + ), + ) + assert "layer.1.k_norm" not in weights + + assert "layer.0.w_gate" in weights + assert "layer.0.w_up" in weights + assert "layer.0.w_down" in weights + assert "layer.1.w_gate" not in weights + assert "layer.1.w_up" not in weights + assert "layer.1.w_down" not in weights + + np.testing.assert_allclose(weights["final_norm"], np.ones(8, dtype=np.float32)) + np.testing.assert_allclose( + weights["w_lm_head"], tensors["model.embed_tokens.weight"].T.astype(np.float32) + ) + + assert weights["_layer_types"] == ["deltanet", "attention"] + assert weights["_num_mamba_layers"] == 1 + assert weights["_num_attention_layers"] == 1 + assert weights["_partial_rotary_factor"] == 0.5 + assert weights["_rope_theta"] == 321.0 + + +def test_get_bundle_config_overrides_normalizes_hybrid_fields(): + """Intent: validate bundle-config normalization for hybrid runtime fields. + Preconditions: text_config provides aliased layer types and linear dimensions. + Postconditions: output reports normalized layer types, counts, and derived dims. + """ + raw = { + "text_config": { + "layer_types": ["linear_attention", "full_attention", "unknown"], + "linear_num_value_heads": 3, + "linear_num_key_heads": 1, + "linear_value_head_dim": 4, + "linear_conv_kernel_dim": 5, + } + } + cfg = ModelConfig( + model_type="qwen3_8", + vocab_size=5, + hidden_size=12, + intermediate_size=16, + num_hidden_layers=3, + num_attention_heads=3, + num_key_value_heads=1, + raw=raw, + ) + + overrides = qwen3_8.plugin.get_bundle_config_overrides(cfg) + assert overrides["layer_types"] == ["deltanet", "attention", "unknown"] + assert overrides["num_mamba_layers"] == 1 + assert overrides["num_attention_layers"] == 1 + assert overrides["d_inner"] == 12 + assert overrides["mamba_d_state"] == 4 + assert overrides["mamba_d_conv"] == 5 + assert overrides["mamba_nheads"] == 3 + assert overrides["mamba_head_dim"] == 4 + assert overrides["conv_dim"] == 20 + + +def _config_from_raw(raw: dict) -> ModelConfig: + text_cfg = raw.get("text_config", {}) + return ModelConfig( + model_type=raw.get("model_type", ""), + architectures=raw.get("architectures", []), + hidden_size=text_cfg.get("hidden_size", 8), + num_hidden_layers=text_cfg.get("num_hidden_layers", 2), + num_attention_heads=text_cfg.get("num_attention_heads", 2), + raw=raw, + ) + + +_QWEN38_RAW = { + "model_type": "qwen3_5", + "architectures": ["Qwen3_5ForConditionalGeneration"], + "text_config": { + "hidden_size": 5120, + "num_hidden_layers": 64, + "num_attention_heads": 24, + "num_key_value_heads": 4, + "head_dim": 256, + "output_gate_type": "swish", + "layer_types": ["linear_attention"] * 3 + ["full_attention"], + }, +} + +_QWEN35_RAW = { + "model_type": "qwen3_5", + "architectures": ["Qwen3_5ForConditionalGeneration"], + "text_config": { + "hidden_size": 4096, + "num_hidden_layers": 32, + "num_attention_heads": 16, + "num_key_value_heads": 2, + "head_dim": 256, + "mlp_only_layers": [], + "layer_types": ["linear_attention"] * 3 + ["full_attention"], + }, +} + + +def test_matches_config_claims_qwen38_and_releases_qwen35(): + """Qwen3.8 ships Qwen3.5's model_type and architecture strings verbatim. + + Intent: the config body, not the checkpoint strings, decides ownership. + Preconditions: two configs identical in model_type/architectures, differing + only by the Qwen3.8 `output_gate_type` marker and the Qwen3.5 + `mlp_only_layers` marker. + Postconditions: this family claims only the Qwen3.8 config, so a genuine + Qwen3.5 checkpoint stays free to fall through to the qwen3_5 family. + """ + assert qwen3_8.plugin.matches_config(_config_from_raw(_QWEN38_RAW)) is True + assert qwen3_8.plugin.matches_config(_config_from_raw(_QWEN35_RAW)) is False + + +def test_matches_config_rejects_other_qwen38_architectures(): + """The dense family must not claim its MoE or qwen4_exp siblings. + + Qwen3.8-2.4T-A95B and Qwen3.8-Flash-Next both carry `output_gate_type`, so + the marker alone is not sufficient -- the architecture/model_type gate is + what keeps them out. + """ + moe = { + "model_type": "qwen3_5_moe_text", + "architectures": ["Qwen3_5MoeForCausalLM"], + "text_config": {"output_gate_type": "swish"}, + } + flash_next = { + "model_type": "qwen4_exp", + "architectures": ["Qwen4ExpForConditionalGeneration"], + "text_config": {"output_gate_type": "sigmoid"}, + } + assert qwen3_8.plugin.matches_config(_config_from_raw(moe)) is False + assert qwen3_8.plugin.matches_config(_config_from_raw(flash_next)) is False + + +def test_matches_by_id_does_not_claim_qwen35_strings(): + assert qwen3_8.plugin.matches("qwen3_8") is True + assert qwen3_8.plugin.matches("qwen3.8") is True + assert qwen3_8.plugin.matches("qwen3_5") is False + + +def test_bundle_overrides_publish_flat_decoder_dims(): + """The C++ runtime reads bundle config with a top-level nlohmann lookup. + + Intent: `hidden_size`/`num_attention_heads`/`num_key_value_heads`/`head_dim` + must appear at the top level of the bundle config, not only under + `text_config`. + Postconditions: without these, `compute_kv_dim()` returns 0 and the KV cache + allocates zero-sized tensors, so pipeline construction fails. + """ + overrides = qwen3_8.plugin.get_bundle_config_overrides( + _config_from_raw(_QWEN38_RAW)) + + assert overrides["hidden_size"] == 5120 + assert overrides["num_attention_heads"] == 24 + assert overrides["num_key_value_heads"] == 4 + assert overrides["head_dim"] == 256 + assert overrides["num_hidden_layers"] == 64 + # kv_dim as the runtime computes it: num_key_value_heads * head_dim. + assert overrides["num_key_value_heads"] * overrides["head_dim"] == 1024 + # eos_token_id stays with the builder, which sources the full stop-id list + # from generation_config.json rather than the single text_config value. + assert "eos_token_id" not in overrides + + +def test_mock_bundle_serializes_decoder_and_hybrid_config(tmp_path): + """Exercise Qwen3.8's nested producer contract with mocked engines.""" + from tensorrt_model_connect.engine_builder import build_bundle + + layer_types = [ + "full_attention" if (index + 1) % 4 == 0 else "linear_attention" + for index in range(64) + ] + # Qwen/Qwen3.8-27B at 1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0. + source_config = { + "architectures": ["Qwen3_5ForConditionalGeneration"], + "image_token_id": 248056, + "model_type": "qwen3_5", + "text_config": { + "bos_token_id": 248044, + "eos_token_id": 248044, + "head_dim": 256, + "hidden_size": 5120, + "intermediate_size": 17408, + "layer_types": layer_types, + "linear_conv_kernel_dim": 4, + "linear_key_head_dim": 128, + "linear_num_key_heads": 16, + "linear_num_value_heads": 48, + "linear_value_head_dim": 128, + "max_position_embeddings": 262144, + "model_type": "qwen3_5_text", + "num_attention_heads": 24, + "num_hidden_layers": 64, + "num_key_value_heads": 4, + "output_gate_type": "swish", + "rms_norm_eps": 1e-6, + "rope_parameters": { + "partial_rotary_factor": 0.25, + "rope_theta": 10000000, + }, + "vocab_size": 248320, + }, + } + (tmp_path / "config.json").write_text( + json.dumps(source_config), + encoding="utf-8", + ) + # Qwen3.8 terminates on 248046, which appears only here; text_config carries + # the single id 248044. + (tmp_path / "generation_config.json").write_text( + json.dumps({"bos_token_id": 248044, "eos_token_id": [248046, 248044]}), + encoding="utf-8", + ) + + class MockQwen38Plugin: + name = "qwen3_8" + runtime_strategy = "qwen3_8_hybrid_mamba_attention" + requires_tokenizer = False + + @staticmethod + def load_weights(_model_dir, _config): + return {} + + @staticmethod + def build_engine(_config, _weights, _max_cache_length, **_kwargs): + return b"MOCK_HYBRID_PLAN" + + @staticmethod + def get_bundle_config_overrides(config): + return qwen3_8.plugin.get_bundle_config_overrides(config) + + with ( + patch( + "tensorrt_model_connect.engine_builder.find_plugin", + return_value=MockQwen38Plugin(), + ), + patch( + "tensorrt_model_connect.engine_builder._get_trt_version", + return_value="11.1.0", + ), + patch( + "tensorrt_model_connect.engine_builder._get_gpu_name", + return_value="CPU unit mock", + ), + patch("tensorrt_model_connect.engine_builder.write_bundle") as write_bundle, + ): + build_bundle( + str(tmp_path), + str(tmp_path / "qwen38-27b.bundle"), + max_cache_length=256, + ) + + sections = { + section.name: section.data for section in write_bundle.call_args.args[2] + } + runtime_config = json.loads(sections["config.json"]) + + # text_config survives untouched for the Python side. + assert runtime_config["text_config"] == source_config["text_config"] + + # The flat decoder contract the strict C++ parser reads. None of these keys + # exist at the top level of the source config. + decoder_contract = { + "vocab_size": 248320, + "hidden_size": 5120, + "num_hidden_layers": 64, + "num_attention_heads": 24, + "num_key_value_heads": 4, + "head_dim": 256, + "bos_token_id": 248044, + } + assert all(key not in source_config for key in decoder_contract) + assert {key: runtime_config[key] for key in decoder_contract} == decoder_contract + # compute_kv_dim() reads these two; zero here means a zero-sized KV cache. + assert runtime_config["num_key_value_heads"] * runtime_config["head_dim"] == 1024 + + # The divergence from qwen3_5: eos_token_id must NOT be republished as an + # override. Overrides are merged last, so doing so would collapse the + # generation_config list to the single text_config id and leave 248046 + # unmatched, running generation to max_new_tokens. + assert runtime_config["eos_token_id"] == [248046, 248044] + + assert runtime_config["layer_types"] == [ + "attention" if layer_type == "full_attention" else "deltanet" + for layer_type in layer_types + ] + assert runtime_config["num_mamba_layers"] == 48 + assert runtime_config["num_attention_layers"] == 16 + assert runtime_config["d_inner"] == 6144 + assert runtime_config["mamba_d_state"] == 128 + assert runtime_config["mamba_d_conv"] == 4 + assert runtime_config["mamba_nheads"] == 48 + assert runtime_config["mamba_head_dim"] == 128 + assert runtime_config["conv_dim"] == 10240 diff --git a/tests/e2e/models/qwen3_8/test_qwen3_8_schedule.py b/tests/e2e/models/qwen3_8/test_qwen3_8_schedule.py new file mode 100644 index 000000000..5c9f79827 --- /dev/null +++ b/tests/e2e/models/qwen3_8/test_qwen3_8_schedule.py @@ -0,0 +1,24 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Qwen3.8-owned scheduler classification checks.""" + +from __future__ import annotations + +import json +from pathlib import Path + +from tests.e2e_harness.manifest_loader import find_manifest_path +from tools.ci import e2e_schedule as schedule_e2e + + +_REPO_ROOT = Path(__file__).resolve().parents[4] + + +def test_qwen38_is_marked_exclusive_gpu() -> None: + models_dir = _REPO_ROOT / "tests" / "e2e" / "models" + manifest_path = find_manifest_path("qwen38-27b", models_dir) + assert manifest_path is not None + manifest = json.loads(manifest_path.read_text()) + + assert schedule_e2e.classify_parallel_resource(manifest) == "exclusive_gpu" diff --git a/tests/e2e/models/qwen3_8/thresholds/qwen38-27b.json b/tests/e2e/models/qwen3_8/thresholds/qwen38-27b.json new file mode 100644 index 000000000..b47a58f38 --- /dev/null +++ b/tests/e2e/models/qwen3_8/thresholds/qwen38-27b.json @@ -0,0 +1,13 @@ +{ + "threshold_overrides": { + "layer_atol": 0.05, + "logit_atol": 0.005, + "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/e2e/timing_estimates.json b/tests/e2e/timing_estimates.json index 936fd55c0..3ccb71095 100644 --- a/tests/e2e/timing_estimates.json +++ b/tests/e2e/timing_estimates.json @@ -69,6 +69,7 @@ "qwen3-moe-tiny-random": 26, "qwen3-vl-2b": 167, "qwen35-9b": 310, + "qwen38-27b": 900, "riva-translate-4b": 232, "roberta-base": 30, "roberta-large": 46, diff --git a/tests/runtime_strategy_matrix.yaml b/tests/runtime_strategy_matrix.yaml index 7d2a07fbb..69999fe4d 100644 --- a/tests/runtime_strategy_matrix.yaml +++ b/tests/runtime_strategy_matrix.yaml @@ -30,6 +30,7 @@ "rwkv_recurrent", "nemotron_h_hybrid_mamba_attention", "qwen3_5_hybrid_mamba_attention", + "qwen3_8_hybrid_mamba_attention", "lfm2_hybrid_conv_attention", "albert_encoder_only", "bert_encoder_only", @@ -762,6 +763,17 @@ "diff_framework_exemption": "No diff_framework check currently registers runtime_strategies=['qwen3_5_hybrid_mamba_attention'].", "performance_mode": "decode" }, + "qwen3_8_hybrid_mamba_attention": { + "task_strategy": "text_generation_causal", + "cli_commands": [ + "run" + ], + "runner_class": "text_generation.TextGenerationCausalRunner", + "comparator_class": "text.TextComparator", + "diff_framework_check_classes": [], + "diff_framework_exemption": "No diff_framework check currently registers runtime_strategies=['qwen3_8_hybrid_mamba_attention'].", + "performance_mode": "decode" + }, "lfm2_hybrid_conv_attention": { "task_strategy": "text_generation_causal", "cli_commands": [ diff --git a/tests/tools/test_family_specialization.py b/tests/tools/test_family_specialization.py index afbaa6dea..e8005b878 100644 --- a/tests/tools/test_family_specialization.py +++ b/tests/tools/test_family_specialization.py @@ -404,7 +404,7 @@ def test_repository_registers_all_current_families() -> None: families = specialization.family_dirs(repo_root, ()) - assert len(families) == 86 + assert len(families) == 87 assert any(family.name == "cosmos3" for family in families) assert any(family.name == "dinov3" for family in families) assert any(family.name == "fast_foundation_stereo" for family in families) diff --git a/tests/tools/test_perf_matrix.py b/tests/tools/test_perf_matrix.py index c12320797..fda77b8a0 100644 --- a/tests/tools/test_perf_matrix.py +++ b/tests/tools/test_perf_matrix.py @@ -270,8 +270,8 @@ def test_release_suite_covers_every_non_l0_ready_model_profile() -> None: performance_catalog.validate_release_coverage(cases, excluded_profiles) - assert len(cases) == 110 - assert len(raw_entries) == 80 + assert len(cases) == 111 + assert len(raw_entries) == 81 assert len(raw_additional) == 30 assert excluded_profiles == { "lfm2-1.2b": LFM2_EXCLUSION_REASON, @@ -289,15 +289,15 @@ def test_release_suite_covers_every_non_l0_ready_model_profile() -> None: assert not any("priority" in entry for entry in raw_entries) assert {case["model"] for case in cases} == ready_profiles - set(excluded_profiles) assert not any(performance_catalog.is_l0_profile(case["model"]) for case in cases) - assert len({(case["family"], case["operation"]) for case in cases}) == 80 - assert len({case["family"] for case in cases}) == 78 + assert len({(case["family"], case["operation"]) for case in cases}) == 81 + assert len({case["family"] for case in cases}) == 79 assert [case["operation"] for case in cases if case["family"] == "eagle_vlm"] == [ "embed", "rerank", ] assert Counter(perf_matrix._candidate_timing_scope(case) for case in cases) == { "model_call_wall": 25, - "public_pipeline_call_wall": 85, + "public_pipeline_call_wall": 86, } assert {case["id"] for case in cases if case["baseline"]["asset_loading_included"]} == { "canary.transcribe", @@ -2191,7 +2191,7 @@ def preflight_after_pending_report(cases, options): assert not scratch_root.exists() results = json.loads((output / "results.json").read_text(encoding="utf-8")) rows = {row["id"]: row for row in results["cases"]} - assert len(rows) == 110 + assert len(rows) == 111 assert results["environment_config"]["name"] == "test-gb300" assert results["environment_config"]["execution"]["minimum_gpu_free_fraction"] == 0.0 assert results["environment_config"]["source"] == str(environment.resolve()) diff --git a/tests/tools/test_performance_catalog.py b/tests/tools/test_performance_catalog.py index b407da1e1..873620538 100644 --- a/tests/tools/test_performance_catalog.py +++ b/tests/tools/test_performance_catalog.py @@ -18,7 +18,7 @@ def test_release_suite_loads_and_selects_models_in_request_order() -> None: selected = suite.select(models=["distilgpt2", "gpt2-125m"]) assert [case["model"] for case in selected] == ["distilgpt2", "gpt2-125m"] - assert len(suite.cases) == 110 + assert len(suite.cases) == 111 def test_release_suite_includes_fast_foundation_stereo() -> None: diff --git a/tests/tools/test_trtmc_validate.py b/tests/tools/test_trtmc_validate.py index 59f543b05..6b8780407 100644 --- a/tests/tools/test_trtmc_validate.py +++ b/tests/tools/test_trtmc_validate.py @@ -45,7 +45,7 @@ def test_model_workload_catalog_covers_every_ready_model(): task_models=task_models, ) - assert len(catalog["models"]) == len(ready_models) == 118 + assert len(catalog["models"]) == len(ready_models) == 119 assert sum("not_compared_reason" in spec for spec in catalog["models"].values()) == 0 assert all("e2e" not in spec.get("workloads", []) for spec in catalog["models"].values()) assert "reference_cache_identity" not in catalog["models"]["personaplex-7b"] @@ -64,7 +64,7 @@ def test_model_workload_catalog_covers_every_ready_model(): } assert len(qwen_identities) == 1 bindings = trtmc_validate.resolve_bindings(catalog, catalog["models"]) - assert len(bindings) == 119 + assert len(bindings) == 120 assert { binding.model for binding in bindings if binding.workload == "mmlu_continuation_parity" } >= { @@ -256,7 +256,7 @@ def test_every_dataset_backed_validation_binding_has_native_reference_runner(): missing.append((model_name, workload, dataset_kind)) assert not missing - assert len({model for model, _workload in bindings}) == 118 + assert len({model for model, _workload in bindings}) == 119 def test_shadow_gate_metrics_include_plugin_task_accuracy() -> None: diff --git a/tests/validation/model_workloads.yaml b/tests/validation/model_workloads.yaml index d25c4f4a9..c7f05b61b 100644 --- a/tests/validation/model_workloads.yaml +++ b/tests/validation/model_workloads.yaml @@ -255,6 +255,8 @@ models: workloads: [vlm_mmmu_pro_vision_fixed_mcq] qwen35-9b: workloads: [mmlu_five_shot_mcq] + qwen38-27b: + workloads: [mmlu_five_shot_mcq] riva-translate-4b: workloads: [flores200_en_fr_riva_translation_parity] roberta-base: diff --git a/tests/validation/workloads.yaml b/tests/validation/workloads.yaml index 231a35715..19838297c 100644 --- a/tests/validation/workloads.yaml +++ b/tests/validation/workloads.yaml @@ -26,6 +26,7 @@ suites: - qwen3-0.6b-topp - qwen3-4b-instruct-2507 - qwen35-9b + - qwen38-27b - qwen3-moe-30b-a3b dataset: kind: mmlu_five_shot_json diff --git a/tools/legal_header_exceptions.toml b/tools/legal_header_exceptions.toml index d3a53efd8..abcd14b9e 100644 --- a/tools/legal_header_exceptions.toml +++ b/tools/legal_header_exceptions.toml @@ -29,4 +29,4 @@ path = "tests/runtime_strategy_matrix.yaml" reason = "JSON-formatted test data is consumed by strict JSON parsers; comments would change parse behavior." license = "Apache-2.0" source = "https://github.com/NVIDIA/TensorRT-Model-Connect/blob/76c8164d97e645a996a210219ba635c8ec9a3453/tests/runtime_strategy_matrix.yaml" -sha256 = "d574a8cd38d235eee8c9a3863b1dcc8160684352b7b3b3b9ab77bced73e3f15b" +sha256 = "f46dc2582ecf6bd725d07ed93fc6a12b1aaac184fe4de6b75e580dc5a0bc0623" diff --git a/website/data/hf-model-metadata.json b/website/data/hf-model-metadata.json index d25bc32c2..b2bc88e62 100644 --- a/website/data/hf-model-metadata.json +++ b/website/data/hf-model-metadata.json @@ -300,6 +300,17 @@ ], "architecture_source": "config.architectures" }, + { + "hf_id": "Qwen/Qwen3.8-27B", + "revision": "1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0", + "revision_source": "resolved", + "metadata_file": "config.json", + "model_type": "qwen3_5", + "architectures": [ + "Qwen3_5ForConditionalGeneration" + ], + "architecture_source": "config.architectures" + }, { "hf_id": "RWKV/rwkv-4-169m-pile", "revision": "46bdc280eb97b6141d5d51a935e0c4870ecaefcc", diff --git a/website/data/model-support-matrix.md b/website/data/model-support-matrix.md index 66494a98a..806a7236c 100644 --- a/website/data/model-support-matrix.md +++ b/website/data/model-support-matrix.md @@ -90,6 +90,7 @@ SPDX-License-Identifier: Apache-2.0 | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | `qwen3-omni-30b-a3b-instruct` | `BF16` | None | — | 🔴 Red | | `Qwen/Qwen3-VL-2B-Instruct` | `qwen3-vl-2b` | `FP16`
FP32 layers: `0, 1, 2` | None | TensorRT Edge-LLM
Qualified TRTMC dispatch target: Coming soon | 🟢 Green | | `Qwen/Qwen3.5-9B` | `qwen35-9b` | `FP16` | None | TensorRT Edge-LLM
Qualified TRTMC dispatch target: Coming soon | 🟢 Green | +| `Qwen/Qwen3.8-27B` | `qwen38-27b` | `FP16` | None | — | 🟡 Yellow | | `nvidia/Riva-Translate-4B-Instruct-v1.1` | `riva-translate-4b` | `FP16` | None | — | 🟢 Green | | `FacebookAI/roberta-base` | `roberta-base` | `FP16` | None | — | 🟢 Green | | `FacebookAI/roberta-large` | `roberta-large` | `FP16` | None | — | 🟢 Green | diff --git a/website/docs/features/model-families.md b/website/docs/features/model-families.md index cd7b591a2..2e2c5429b 100644 --- a/website/docs/features/model-families.md +++ b/website/docs/features/model-families.md @@ -33,7 +33,7 @@ Common native TensorRT groups: - Decoder-only: Qwen, LLaMA, Mistral, GPT, OPT, Bloom, Gemma, Falcon, Granite, OLMo. - MoE: Mixtral, Phi-MoE, Qwen-MoE, GPT-OSS, DeepSeek-V2. -- Recurrent and hybrid: Mamba, RWKV, Nemotron-H, Qwen3.5 hybrid. +- Recurrent and hybrid: Mamba, RWKV, Nemotron-H, Qwen3.5 hybrid, Qwen3.8 hybrid. - Encoder-only: BERT, RoBERTa, DeBERTa, ModernBERT, DistilBERT, ConvBERT, FNet, XLNet, MPNet, DPR. - Seq2seq: T5, Marian, BART, M2M-100. - Vision-language: Qwen-VL, InternVL, Lance, LocateAnything, Phi4 multimodal, diff --git a/website/docs/features/runtime-strategies.md b/website/docs/features/runtime-strategies.md index d17e8c57b..7bbaab77b 100644 --- a/website/docs/features/runtime-strategies.md +++ b/website/docs/features/runtime-strategies.md @@ -17,7 +17,7 @@ selects that path before reading a native strategy. | Category | Representative model-owned strategies | | --- | --- | | Text decoder | `qwen_decoder_kv_cache`, `llama_decoder_kv_cache`, `mixtral_decoder_moe`, `gpt_oss_decoder_moe` | -| Recurrent text | `mamba_ssm_recurrent`, `rwkv_recurrent`, `nemotron_h_hybrid_mamba_attention`, `qwen3_5_hybrid_mamba_attention` | +| Recurrent text | `mamba_ssm_recurrent`, `rwkv_recurrent`, `nemotron_h_hybrid_mamba_attention`, `qwen3_5_hybrid_mamba_attention`, `qwen3_8_hybrid_mamba_attention` | | Encoder and retrieval | `bert_encoder_only`, `mpnet_encoder_only`, `eagle_vlm_embedding`, `eagle_vlm_reranking` | | Seq2seq | `t5_text_to_text`, `marian_translation`, `bart_seq2seq_encoder_decoder`, `m2m_100_seq2seq_encoder_decoder` | | Vision and multimodal | `qwen_vl_vision_language`, `internvl_vision_language`, `qwen3_omni_multimodal` |