[Perf][MiniCPM-o] TensorRT acceleration for the Code2Wav vocoder (DiT estimator + campplus) (#5638)

Signed-off-by: Yuekai Zhang <zhangyuekai@foxmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>

pyproject.toml

@@ -70,8 +70,7 @@ dev = [
     "FlagEmbedding==1.4.0", # for hunyuan_image3 tests
     "kernels==0.14.1", # for accuracy tests comparison with diffusers output (0.15.0 incompatible with diffusers 0.38.0)
     "s3tokenizer==0.3.0", # for tts test
-    "step-audio2==1.0.0", # Step-Audio2 Token2Wav (flashcosyvoice / hyperpyyaml)
-    "stepaudio2-minicpmo", # for MiniCPM-o 4.5 talker vocoder (stepaudio2-minicpmo Token2wav)
+    "step-audio2==1.0.0", # Step-Audio2 / MiniCPM-o 4.5 Token2Wav (flashcosyvoice / hyperpyyaml)
 ]
 
 demo = [

vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py

@@ -60,9 +60,12 @@ class BatchedToken2Wav(nn.Module):
     asset loader and prompt feature extractor.
     """
 
-    def __init__(self, token2wav: Any):
+    def __init__(self, token2wav: Any, trt_stepper: Any | None = None):
         super().__init__()
         self._token2wav = token2wav
+        # Optional TrtDiTStepper (step_audio2_dit_trt): replaces only the
+        # per-timestep DiT estimator call; encoder and HiFT stay on torch.
+        self._trt_stepper = trt_stepper
         self.flow = token2wav.flow
         self.hift = token2wav.hift
         hift_parameter = next(self.hift.parameters(), None)
@@ -200,6 +203,17 @@ class BatchedToken2Wav(nn.Module):
         cnn_cache: torch.Tensor | None,
         att_cache: torch.Tensor | None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if self._trt_stepper is not None:
+            out, new_cnn, new_att = self._trt_stepper.step(
+                x=x,
+                mu=mu,
+                t=time,
+                spks=speakers,
+                cond=cond,
+                cnn_cache=cnn_cache,
+                att_cache=att_cache,
+            )
+            return out.to(mu.dtype), new_cnn, new_att
         time_embedding = estimator.t_embedder(time).unsqueeze(1)
         width = int(x.shape[-1])
         speaker_features = speakers.unsqueeze(-1).expand(-1, -1, width)

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py

@@ -727,19 +727,17 @@ class MiniCPMO45Code2Wav(nn.Module):
         if self.backend is not None:
             return
 
+        # In-tree adapter over StepAudio2Token2WavCore on every platform. It
+        # matches the external `stepaudio2-minicpmo` Token2wav bit-for-bit on
+        # CUDA (same cosyvoice2 flow/DiT modules and weights) while dropping
+        # that package's hard-coded `.cuda()` calls, and auto-applies the
+        # Ascend fixes (HiFT linear downsample, DiT mask expand, MATH SDPA)
+        # on NPU.
+        from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_token2wav import (
+            MiniCPMO45Token2wav as Token2wav,
+        )
         from vllm_omni.platforms import current_omni_platform
 
-        if current_omni_platform.is_npu():
-            # NPU/Ascend: the external `stepaudio2` package hard-codes `.cuda()`,
-            # so use the in-tree NPU-aware adapter instead. It delegates to
-            # StepAudio2Token2WavCore, which auto-applies the Ascend fixes
-            # (HiFT linear downsample, DiT mask expand, MATH SDPA) on NPU.
-            from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_token2wav import (
-                MiniCPMO45Token2wav as Token2wav,
-            )
-        else:
-            from stepaudio2.token2wav import Token2wav
-
         extra = self._extra_config()
         # Hub repo ids only need to become local directories once the vocoder
         # assets are actually read; unit tests construct this model with fake
@@ -765,4 +763,26 @@ class MiniCPMO45Code2Wav(nn.Module):
             )
         finally:
             torch.set_default_dtype(previous_dtype)
-        self.backend = BatchedToken2Wav(token2wav)
+
+        trt_stepper = None
+        use_trt = bool(extra.get("token2wav_trt", False)) or os.environ.get("MINICPMO_TOKEN2WAV_TRT", "") == "1"
+        # TensorRT is CUDA-only; other platforms ignore the toggle.
+        if use_trt and current_omni_platform.is_cuda():
+            from vllm_omni.model_executor.models.step_audio2.step_audio2_dit_trt import build_dit_trt_stepper
+
+            dtype_name = str(
+                extra.get("token2wav_trt_dtype", os.environ.get("MINICPMO_TOKEN2WAV_TRT_DTYPE", "fp16"))
+            ).lower()
+            trt_dtype = torch.float32 if dtype_name in ("fp32", "float32") else torch.float16
+            max_batch = int(extra.get("token2wav_trt_max_batch", 16))
+            device = next(token2wav.flow.parameters()).device
+            trt_stepper = build_dit_trt_stepper(
+                token2wav.flow.decoder.estimator,
+                device=device,
+                dtype=trt_dtype,
+                max_batch=max_batch,
+            )
+            logger.info("MiniCPM-o Code2Wav: DiT estimator running on TensorRT (%s)", dtype_name)
+            token2wav.enable_trt_spk_embedding()
+            logger.info("MiniCPM-o Code2Wav: campplus speaker embedding running on TensorRT")
+        self.backend = BatchedToken2Wav(token2wav, trt_stepper=trt_stepper)

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_token2wav.py

@@ -120,6 +120,10 @@ class MiniCPMO45Token2wav:
         """Delegate prompt feature extraction to the wrapped core."""
         return self._core._prepare_prompt(prompt_wav)
 
+    def enable_trt_spk_embedding(self) -> None:
+        """Run the campplus speaker-embedding model on TensorRT."""
+        self._core.enable_trt_spk_embedding()
+
     def __call__(self, generated_speech_tokens, prompt_wav) -> bytes:
         """One-shot tokens → 24 kHz WAV bytes."""
         return self._core.forward(

vllm_omni/model_executor/models/step_audio2/step_audio2_dit_trt.py

@@ -0,0 +1,375 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""TensorRT acceleration for the Step-Audio2 / CosyVoice2 flow DiT estimator.
+
+Ports stepfun-ai/Step-Audio2#70 into vLLM-Omni. The upstream PR patches
+``cosyvoice2.flow.flow_matching`` to dispatch between a torch estimator and a
+raw TRT execution context; here the swap happens one level up instead — the
+batched Code2Wav driver (``BatchedToken2Wav._estimator_step``) already owns
+the per-timestep estimator call, so a :class:`TrtDiTStepper` replaces only
+that call and the pip-installed ``cosyvoice2`` package stays untouched.
+
+Differences from the upstream PR:
+
+* ONNX export is a function (:func:`export_dit_chunk_onnx`) that wraps the
+  already-loaded torch ``DiT`` module and runs at engine-build time — no
+  separate ``tools/export_onnx_*.py`` step.
+* The engine executes on the caller's current CUDA stream, so no context
+  pool, no private stream, and no ``synchronize()`` per call (a private
+  stream without synchronization produces silent garbage audio; with
+  synchronization it costs two round-trips per denoise step).
+* Cache tensors are exchanged as stacked ``(depth, 2B, ...)`` tensors, the
+  same layout ``BatchedToken2Wav`` already uses.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_TRT_CACHE_ENV = "MINICPMO_TOKEN2WAV_TRT_CACHE"
+
+
+class _DiTChunkWrapper(nn.Module):
+    """ONNX-exportable streaming ``forward_chunk`` with explicit stacked caches.
+
+    Mirrors ``DiT.forward_chunk`` (cosyvoice2 pip package) but takes and
+    returns caches as plain tensors instead of registered buffers, matching
+    the I/O contract of the TRT engine and of ``BatchedToken2Wav``.
+    """
+
+    def __init__(self, dit: nn.Module):
+        super().__init__()
+        self.dit = dit
+
+    def forward(
+        self,
+        x: torch.Tensor,  # (2B, mel, T)
+        mu: torch.Tensor,  # (2B, mel, T)
+        t: torch.Tensor,  # (2B,)
+        spks: torch.Tensor,  # (2B, spk)
+        cond: torch.Tensor,  # (2B, mel, T)
+        cnn_cache: torch.Tensor,  # (depth, 2B, c1+c2, kernel-1)
+        att_cache: torch.Tensor,  # (depth, 2B, heads, L, head_dim*2)
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        t_emb = self.dit.t_embedder(t).unsqueeze(1)
+        spks_f = spks.unsqueeze(-1).expand(-1, -1, x.shape[-1])
+        h = torch.cat((x, mu, spks_f, cond), dim=1).transpose(1, 2)
+        h = self.dit.in_proj(h)
+        new_cnn = []
+        new_att = []
+        for i, block in enumerate(self.dit.blocks):
+            h, block_cnn, block_att = block.forward_chunk(h, t_emb, cnn_cache[i], att_cache[i], None)
+            new_cnn.append(block_cnn)
+            new_att.append(block_att)
+        h = self.dit.final_layer(h, t_emb)
+        return h.transpose(1, 2), torch.stack(new_cnn), torch.stack(new_att)
+
+
+def _dit_dims(dit: nn.Module) -> dict[str, int]:
+    block0 = dit.blocks[0]
+    return {
+        "depth": len(dit.blocks),
+        "mel": int(dit.out_channels),
+        "spk": int(dit.in_channels) - 3 * int(dit.out_channels),
+        "heads": int(block0.attn.num_heads),
+        "att_dim": int(block0.attn.head_dim) * 2,
+        "cnn_ch": int(block0.conv.in_channels + block0.conv.out_channels),
+        "cnn_w": int(block0.conv.block[1].causal_padding[0]),
+    }
+
+
+def export_dit_chunk_onnx(
+    dit: nn.Module,
+    onnx_path: str | Path,
+    device: torch.device,
+    dtype: torch.dtype = torch.float16,
+) -> None:
+    """Export the streaming DiT chunk graph to ONNX in ``dtype``.
+
+    Fused-in replacement for the upstream ``tools/export_onnx_streaming_token2wav.py``.
+    TensorRT >= 10 deprecates weakly-typed networks (TRT 11 removes
+    ``BuilderFlag.FP16`` outright), so the engine precision comes from the
+    ONNX graph itself. The DiT is deep-copied for the cast — the live torch
+    estimator keeps its original dtype.
+    """
+    import copy
+
+    dims = _dit_dims(dit)
+    export_dit = copy.deepcopy(dit).to(device=device, dtype=dtype).eval()
+    wrapper = _DiTChunkWrapper(export_dit).eval()
+    batch, chunk, cache_len = 2, 48, 16
+    sample = (
+        torch.randn(batch, dims["mel"], chunk, device=device, dtype=dtype),
+        torch.randn(batch, dims["mel"], chunk, device=device, dtype=dtype),
+        torch.rand(batch, device=device, dtype=dtype),
+        torch.randn(batch, dims["spk"], device=device, dtype=dtype),
+        torch.randn(batch, dims["mel"], chunk, device=device, dtype=dtype),
+        torch.zeros(dims["depth"], batch, dims["cnn_ch"], dims["cnn_w"], device=device, dtype=dtype),
+        torch.randn(dims["depth"], batch, dims["heads"], cache_len, dims["att_dim"], device=device, dtype=dtype),
+    )
+    input_names = ["x", "mu", "t", "spks", "cond", "cnn_cache", "att_cache"]
+    output_names = ["out", "new_cnn_cache", "new_att_cache"]
+    dynamic_axes = {
+        "x": {0: "b", 2: "t"},
+        "mu": {0: "b", 2: "t"},
+        "t": {0: "b"},
+        "spks": {0: "b"},
+        "cond": {0: "b", 2: "t"},
+        "cnn_cache": {1: "b"},
+        "att_cache": {1: "b", 3: "l"},
+        "out": {0: "b", 2: "t"},
+        "new_cnn_cache": {1: "b"},
+        "new_att_cache": {1: "b", 3: "l_out"},
+    }
+    logger.info("Exporting DiT chunk ONNX to %s", onnx_path)
+    onnx_path = Path(onnx_path)
+    onnx_path.parent.mkdir(parents=True, exist_ok=True)
+    # Process-unique temp name + atomic rename, so concurrent first builds on
+    # a shared cache dir never see (or clobber) a partially written file.
+    tmp = onnx_path.with_suffix(f".tmp.{os.getpid()}")
+    try:
+        with torch.no_grad():
+            torch.onnx.export(
+                wrapper,
+                sample,
+                str(tmp),
+                input_names=input_names,
+                output_names=output_names,
+                dynamic_axes=dynamic_axes,
+                opset_version=17,
+                dynamo=False,
+            )
+        os.replace(tmp, onnx_path)
+    finally:
+        tmp.unlink(missing_ok=True)
+
+
+def convert_onnx_to_trt(
+    onnx_path: str | Path,
+    plan_path: str | Path,
+    dims: dict[str, int],
+    *,
+    max_batch: int = 16,
+    max_len: int = 3000,
+    max_cache_len: int = 1000,
+) -> None:
+    """Build a TRT engine from the exported ONNX (fused-in from upstream tools).
+
+    Engine precision is strongly typed from the ONNX graph itself — export
+    the graph in the desired dtype (TRT 11 removed ``BuilderFlag.FP16``).
+    """
+    import tensorrt as trt
+
+    logger.info("Building DiT TRT engine at %s — this can take a few minutes", plan_path)
+    trt_logger = trt.Logger(trt.Logger.INFO)
+    builder = trt.Builder(trt_logger)
+    # TRT >= 10 removed EXPLICIT_BATCH (networks are always explicit-batch).
+    flag = getattr(trt.NetworkDefinitionCreationFlag, "EXPLICIT_BATCH", None)
+    network = builder.create_network(1 << int(flag) if flag is not None else 0)
+    parser = trt.OnnxParser(network, trt_logger)
+    with open(onnx_path, "rb") as f:
+        if not parser.parse(f.read()):
+            errors = [str(parser.get_error(i)) for i in range(parser.num_errors)]
+            raise ValueError(f"failed to parse {onnx_path}: {errors}")
+
+    config = builder.create_builder_config()
+
+    depth, mel, spk = dims["depth"], dims["mel"], dims["spk"]
+    heads, att_dim = dims["heads"], dims["att_dim"]
+    cnn_ch, cnn_w = dims["cnn_ch"], dims["cnn_w"]
+    opt_b = min(4, max_batch)
+    shapes = {
+        "x": [(2, mel, 2), (opt_b, mel, 500), (max_batch, mel, max_len)],
+        "mu": [(2, mel, 2), (opt_b, mel, 500), (max_batch, mel, max_len)],
+        "t": [(2,), (opt_b,), (max_batch,)],
+        "spks": [(2, spk), (opt_b, spk), (max_batch, spk)],
+        "cond": [(2, mel, 2), (opt_b, mel, 500), (max_batch, mel, max_len)],
+        "cnn_cache": [
+            (depth, 2, cnn_ch, cnn_w),
+            (depth, opt_b, cnn_ch, cnn_w),
+            (depth, max_batch, cnn_ch, cnn_w),
+        ],
+        "att_cache": [
+            (depth, 2, heads, 0, att_dim),
+            (depth, opt_b, heads, 500, att_dim),
+            (depth, max_batch, heads, max_cache_len, att_dim),
+        ],
+    }
+    profile = builder.create_optimization_profile()
+    for name, (mn, op, mx) in shapes.items():
+        profile.set_shape(name, mn, op, mx)
+    config.add_optimization_profile(profile)
+
+    engine_bytes = builder.build_serialized_network(network, config)
+    if engine_bytes is None:
+        raise RuntimeError("TensorRT engine build failed")
+    plan_path = Path(plan_path)
+    plan_path.parent.mkdir(parents=True, exist_ok=True)
+    # Process-unique temp name: concurrent first builds (multiple ranks on a
+    # shared cache dir) must not clobber each other's partial writes; the
+    # final os.replace stays atomic and last-writer-wins with equal content.
+    tmp = plan_path.with_suffix(f".tmp.{os.getpid()}")
+    try:
+        with open(tmp, "wb") as f:
+            f.write(engine_bytes)
+        os.replace(tmp, plan_path)
+    finally:
+        tmp.unlink(missing_ok=True)
+    logger.info("DiT TRT engine written to %s", plan_path)
+
+
+class TrtDiTStepper:
+    """Drop-in TRT replacement for one DiT denoise step.
+
+    ``step`` matches the tensors ``BatchedToken2Wav._decode_cfm`` has in hand
+    per timestep. Executes on the caller's current CUDA stream — same-stream
+    ordering makes explicit synchronization unnecessary.
+    """
+
+    def __init__(
+        self,
+        plan_path: str | Path,
+        dims: dict[str, int],
+        device: torch.device,
+        dtype: torch.dtype,
+        *,
+        max_batch: int = 16,
+        max_len: int = 3000,
+        max_cache_len: int = 1000,
+    ):
+        import tensorrt as trt
+
+        self.dims = dims
+        self.device = device
+        self.dtype = dtype
+        # Optimization-profile bounds the engine was built with; step() checks
+        # them so an out-of-profile call fails with a readable error instead
+        # of an opaque execute_async_v3 failure.
+        self.max_batch = int(max_batch)
+        self.max_len = int(max_len)
+        self.max_cache_len = int(max_cache_len)
+        with open(plan_path, "rb") as f:
+            self.engine = trt.Runtime(trt.Logger(trt.Logger.INFO)).deserialize_cuda_engine(f.read())
+        if self.engine is None:
+            raise RuntimeError(f"failed to deserialize TRT engine {plan_path}")
+        self.context = self.engine.create_execution_context()
+        if self.context is None:
+            raise RuntimeError("failed to create TRT execution context")
+        # TRT rejects null bindings; zero-length cache tensors share this address.
+        self._dummy = torch.zeros(1, device=device, dtype=dtype)
+
+    def _addr(self, tensor: torch.Tensor) -> int:
+        return self._dummy.data_ptr() if tensor.numel() == 0 else tensor.data_ptr()
+
+    def step(
+        self,
+        *,
+        x: torch.Tensor,
+        mu: torch.Tensor,
+        t: torch.Tensor,
+        spks: torch.Tensor,
+        cond: torch.Tensor,
+        cnn_cache: torch.Tensor | None,
+        att_cache: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        dims = self.dims
+        batch = int(x.shape[0])
+        chunk = int(x.shape[2])
+        cache_frames = int(att_cache.shape[3]) if att_cache is not None else 0
+        if batch > self.max_batch or chunk > self.max_len or cache_frames > self.max_cache_len:
+            raise ValueError(
+                f"DiT TRT engine profile exceeded: batch={batch} (max {self.max_batch}, "
+                f"= token2wav_trt_max_batch, CFG included), chunk={chunk} frames "
+                f"(max {self.max_len}), att cache={cache_frames} frames (max {self.max_cache_len}). "
+                "Rebuild the engine with a larger token2wav_trt_max_batch or disable token2wav_trt."
+            )
+        if cnn_cache is None:
+            cnn_cache = torch.zeros(
+                (dims["depth"], batch, dims["cnn_ch"], dims["cnn_w"]), device=self.device, dtype=self.dtype
+            )
+        if att_cache is None:
+            att_cache = torch.zeros(
+                (dims["depth"], batch, dims["heads"], 0, dims["att_dim"]), device=self.device, dtype=self.dtype
+            )
+        cache_len = int(att_cache.shape[3])
+
+        inputs = {
+            "x": x.to(self.dtype).contiguous(),
+            "mu": mu.to(self.dtype).contiguous(),
+            "t": t.to(self.dtype).contiguous(),
+            "spks": spks.to(self.dtype).contiguous(),
+            "cond": cond.to(self.dtype).contiguous(),
+            "cnn_cache": cnn_cache.to(self.dtype).contiguous(),
+            "att_cache": att_cache.to(self.dtype).contiguous(),
+        }
+        out = torch.empty((batch, dims["mel"], chunk), device=self.device, dtype=self.dtype)
+        new_cnn = torch.empty_like(inputs["cnn_cache"])
+        new_att = torch.empty(
+            (dims["depth"], batch, dims["heads"], cache_len + chunk, dims["att_dim"]),
+            device=self.device,
+            dtype=self.dtype,
+        )
+
+        for name, tensor in inputs.items():
+            self.context.set_input_shape(name, tuple(tensor.shape))
+            self.context.set_tensor_address(name, self._addr(tensor))
+        self.context.set_tensor_address("out", out.data_ptr())
+        self.context.set_tensor_address("new_cnn_cache", new_cnn.data_ptr())
+        self.context.set_tensor_address("new_att_cache", self._addr(new_att))
+
+        stream = torch.cuda.current_stream(self.device)
+        if not self.context.execute_async_v3(stream.cuda_stream):
+            raise RuntimeError("TRT DiT execution failed")
+        return out, new_cnn, new_att
+
+
+def _default_cache_dir() -> Path:
+    env = os.environ.get(_TRT_CACHE_ENV, "").strip()
+    if env:
+        return Path(env)
+    hf_home = os.environ.get("HF_HOME", "").strip()
+    base = Path(hf_home) if hf_home else Path.home() / ".cache" / "huggingface"
+    return base / "vllm_omni" / "token2wav_trt"
+
+
+def build_dit_trt_stepper(
+    dit: nn.Module,
+    *,
+    device: torch.device | str = "cuda",
+    dtype: torch.dtype = torch.float16,
+    max_batch: int = 16,
+    cache_dir: str | Path | None = None,
+) -> TrtDiTStepper:
+    """Export (if needed), build (if needed) and load the DiT TRT engine.
+
+    Artifacts are cached under ``cache_dir`` (default: ``$MINICPMO_TOKEN2WAV_TRT_CACHE``
+    or ``$HF_HOME/vllm_omni/token2wav_trt``), keyed by dtype/batch/TRT
+    version/GPU architecture — serialized plans are SM-specific, so a cache
+    shared across heterogeneous machines must not reuse another arch's plan.
+    """
+    import tensorrt as trt
+
+    device = torch.device(device)
+    dims = _dit_dims(dit)
+    cache = Path(cache_dir) if cache_dir is not None else _default_cache_dir()
+    cache.mkdir(parents=True, exist_ok=True)
+    dtype_tag = str(dtype).replace("torch.", "")
+    major, minor = torch.cuda.get_device_capability(device)
+    sm_tag = f"sm{major}{minor}"
+    onnx_path = cache / f"dit_chunk.d{dims['depth']}.{dtype_tag}.onnx"
+    plan_path = cache / (f"dit_chunk.d{dims['depth']}.{dtype_tag}.b{max_batch}.{sm_tag}.trt{trt.__version__}.plan")
+
+    if not plan_path.is_file() or plan_path.stat().st_size == 0:
+        if not onnx_path.is_file() or onnx_path.stat().st_size == 0:
+            export_dit_chunk_onnx(dit, onnx_path, device, dtype)
+        convert_onnx_to_trt(onnx_path, plan_path, dims, max_batch=max_batch)
+    else:
+        logger.info("Reusing cached DiT TRT engine %s", plan_path)
+    return TrtDiTStepper(plan_path, dims, device, dtype, max_batch=max_batch)

vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py

@@ -178,6 +178,28 @@ class StepAudio2Token2WavCore(nn.Module):
         self._ensure_models_loaded()
         return self._hift
 
+    def _forward_spk_embedding(self, spk_feat: torch.Tensor) -> torch.Tensor:
+        """Run campplus on ``[T, 80]`` fbank features; returns ``[1, 192]`` on ``self.device``."""
+        spk_model = self.spk_model
+        if isinstance(spk_model, onnxruntime.InferenceSession):
+            return torch.tensor(
+                spk_model.run(None, {spk_model.get_inputs()[0].name: spk_feat.unsqueeze(dim=0).cpu().numpy()})[0],
+                device=self.device,
+            )
+        # CampplusTRT (or any callable taking [T, 80] and returning [1, 192]).
+        return spk_model(spk_feat).to(self.device)
+
+    def enable_trt_spk_embedding(self) -> None:
+        """Swap the onnxruntime campplus session for the shared TensorRT engine.
+
+        Reuses the CosyVoice3 CAM++ engine wrapper — MiniCPM-o / Step-Audio2
+        ship the same campplus.onnx architecture and I/O contract.
+        """
+        from vllm_omni.model_executor.models.cosyvoice3.speaker_embedding_trt import get_campplus_trt
+
+        self._ensure_models_loaded()
+        self._spk_model = get_campplus_trt(f"{self.model_path}/campplus.onnx", device=self.device)
+
     def _prepare_prompt(self, prompt_wav: str):
         """Prepare prompt audio for conditioning"""
         # Prefer soundfile/librosa path to avoid torchaudio->torchcodec runtime coupling.
@@ -202,10 +224,7 @@ class StepAudio2Token2WavCore(nn.Module):
 
         spk_feat = kaldi.fbank(audio.unsqueeze(0), num_mel_bins=80, dither=0, sample_frequency=16000)
         spk_feat = spk_feat - spk_feat.mean(dim=0, keepdim=True)
-        spk_emb = torch.tensor(
-            self.spk_model.run(None, {self.spk_model.get_inputs()[0].name: spk_feat.unsqueeze(dim=0).cpu().numpy()})[0],
-            device=self.device,
-        )
+        spk_emb = self._forward_spk_embedding(spk_feat)
 
         # 24 kHz branch: mel spectrogram for flow model conditioning
         # Must resample from the ORIGINAL audio, not the 16 kHz version.