"""Synchronous NumPy API over TensorRT, with reusable pinned/device buffers. One instance owns one execution context and must not be called concurrently. CUDA graph mode recaptures whenever the shape changes; use stable length buckets when enabling it in a latency-critical serving loop. """ from __future__ import annotations from pathlib import Path from typing import Mapping import numpy as np import torch import tensorrt as trt class TensorRTBackend: def __init__(self, engine_path: str | Path = "artifacts/model.fp16.engine", device: int = 0, cuda_graph: bool = False, *, _shared=None): if not torch.cuda.is_available(): raise RuntimeError("TensorRT inference requires a CUDA GPU") self.device = torch.device(f"cuda:{device}") self.cuda_graph = cuda_graph self._engine_owner = _shared self.logger = _shared.logger if _shared else trt.Logger(trt.Logger.ERROR) if _shared and _shared.device != self.device: raise ValueError("Shared engine and execution context must use the same GPU") if not _shared: trt.init_libnvinfer_plugins(self.logger, "") self.runtime = _shared.runtime if _shared else trt.Runtime(self.logger) with torch.cuda.device(self.device): self.engine = (_shared.engine if _shared else self.runtime.deserialize_cuda_engine(Path(engine_path).read_bytes())) if self.engine is None: raise RuntimeError(f"Could not deserialize TensorRT engine: {engine_path}") self.context = self.engine.create_execution_context() if self.context is None: raise RuntimeError("Could not create TensorRT execution context") self.context.nvtx_verbosity = trt.ProfilingVerbosity.NONE self.stream = torch.cuda.Stream(device=self.device) self.names = [self.engine.get_tensor_name(i) for i in range(self.engine.num_io_tensors)] self.input_names = [n for n in self.names if self.engine.get_tensor_mode(n) == trt.TensorIOMode.INPUT] self.output_names = [n for n in self.names if self.engine.get_tensor_mode(n) == trt.TensorIOMode.OUTPUT] if len(self.output_names) != 1: raise ValueError(f"Expected one logits output, got {self.output_names}") self.numpy_dtypes = {n: np.dtype(trt.nptype(self.engine.get_tensor_dtype(n))) for n in self.names} self.torch_dtypes = {n: torch.from_numpy(np.empty((), dtype=self.numpy_dtypes[n])).dtype for n in self.names} self._signature = None self._host: dict[str, torch.Tensor] = {} self._host_numpy: dict[str, np.ndarray] = {} self._device: dict[str, torch.Tensor] = {} self._graph = None def _configure(self, arrays: Mapping[str, np.ndarray]) -> None: signature = tuple((n, tuple(arrays[n].shape)) for n in self.input_names) if signature == self._signature: return self.stream.synchronize() self._graph = None for name in self.input_names: if not self.context.set_input_shape(name, arrays[name].shape): minimum, optimum, maximum = self.engine.get_tensor_profile_shape(name, 0) raise ValueError(f"{name} shape {arrays[name].shape} outside engine profile {minimum}..{maximum}") missing = self.context.infer_shapes() if missing: raise RuntimeError(f"Unspecified TensorRT input shapes: {missing}") for name in self.names: shape = tuple(self.context.get_tensor_shape(name)) if any(dim < 0 for dim in shape): raise RuntimeError(f"Unresolved TensorRT output shape {name}: {shape}") self._host[name] = torch.empty(shape, dtype=self.torch_dtypes[name], pin_memory=True) self._host_numpy[name] = self._host[name].numpy() self._device[name] = torch.empty(shape, dtype=self.torch_dtypes[name], device=self.device) if not self.context.set_tensor_address(name, self._device[name].data_ptr()): raise RuntimeError(f"Could not bind TensorRT tensor {name}") self._signature = signature def _enqueue(self) -> None: if not self.context.execute_async_v3(self.stream.cuda_stream): raise RuntimeError("TensorRT execution failed") def run(self, inputs: Mapping[str, np.ndarray]) -> np.ndarray: """Return independent CPU logits; timing includes both H2D and D2H copies.""" missing = set(self.input_names).difference(inputs) if missing: raise ValueError(f"Missing model inputs: {sorted(missing)}") arrays = {n: np.asarray(inputs[n], dtype=self.numpy_dtypes[n]) for n in self.input_names} shapes = [array.shape for array in arrays.values()] if any(len(shape) != 2 or shape != shapes[0] for shape in shapes): raise ValueError(f"All inputs must have the same [batch, sequence] shape; got {shapes}") with torch.cuda.device(self.device), torch.cuda.stream(self.stream): self._configure(arrays) for name in self.input_names: np.copyto(self._host_numpy[name], arrays[name]) self._device[name].copy_(self._host[name], non_blocking=True) if self.cuda_graph: if self._graph is None: # TRT must complete shape-dependent internal updates before # capture. Stable buffers and context remain alive until the # next shape change, which discards the captured graph. self._enqueue() self.stream.synchronize() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph, stream=self.stream): self._enqueue() self._graph = graph self._graph.replay() else: self._enqueue() output = self.output_names[0] self._host[output].copy_(self._device[output], non_blocking=True) self.stream.synchronize() # Callers may retain predictions after subsequent runs. return self._host_numpy[output].copy() # Short alias for projects preferring the name of the runtime. TRTBackend = TensorRTBackend class TensorRTGraphBackend: """Batch-one CUDA graphs retained for five sequence-length buckets. Each bucket owns a fixed execution context and fixed buffers, while all contexts share engine weights. Larger batches use the normal dynamic backend. Calls must be serialized, as with TensorRTBackend. """ def __init__(self, engine_path: str | Path = "artifacts/model.fp16.engine", device: int = 0, pad_token_id: int = 0): self._engine_path = engine_path self._device_index = device self.pad_token_id = pad_token_id self.buckets = (32, 64, 128, 256, 512) self.fallback = TensorRTBackend(engine_path, device=device) self.input_names = self.fallback.input_names self._bucket_backends: dict[int, TensorRTBackend] = {} def run(self, inputs: Mapping[str, np.ndarray]) -> np.ndarray: missing = set(self.input_names).difference(inputs) if missing: raise ValueError(f"Missing model inputs: {sorted(missing)}") arrays = {name: np.asarray(inputs[name]) for name in self.input_names} shapes = [array.shape for array in arrays.values()] if any(len(shape) != 2 or shape != shapes[0] for shape in shapes): raise ValueError(f"All inputs must have the same [batch, sequence] shape; got {shapes}") batch, sequence = shapes[0] if batch != 1: return self.fallback.run(arrays) if not 2 <= sequence <= 512: raise ValueError(f"Sequence length must be in [2, 512]; got {sequence}") bucket = next(length for length in self.buckets if length >= sequence) if bucket not in self._bucket_backends: self._bucket_backends[bucket] = TensorRTBackend( self._engine_path, device=self._device_index, cuda_graph=True, _shared=self.fallback, ) if sequence != bucket: padded = {} for name, array in arrays.items(): value = self.pad_token_id if name == "input_ids" else 0 padded[name] = np.full((1, bucket), value, dtype=array.dtype) padded[name][:, :sequence] = array arrays = padded # TensorRTBackend already returns an independent array; the slice keeps # that storage alive and cannot be overwritten by a subsequent call. return self._bucket_backends[bucket].run(arrays)[:, :sequence]