type: user

role: user

content:
Research whether the following inference frameworks support serving small audio/codec token language models (like VALL-E style TTS that predicts EnCodec tokens):

1. **vLLM** — does it support custom transformer decoder models that predict audio tokens (not text)? Can you load a custom PyTorch transformer model into vLLM? Or is it only for text LLMs?

2. **TensorRT / TensorRT-LLM** — can it optimize custom small transformer models (12M params) that do next-token prediction over a 1024-vocab codebook? Is there support for non-text token prediction?

3. **ONNX Runtime** — support for custom transformer decoders with KV-cache?

4. **CTranslate2** — any audio model support?

5. Any other inference frameworks specifically designed for audio token models?

Also check if any existing TTS systems (VALL-E, XTTS, Parler-TTS, etc.) are deployed via vLLM or TensorRT and how they do it.

Be specific about what works and what doesn't.

userType: external

entrypoint: cli

gitBranch: HEAD

slug: steady-twirling-elephant


type: message

role: assistant

type: text

text:
# Inference Frameworks for Small Audio/Codec Token Language Models

Here is a thorough analysis based on my knowledge of these frameworks (as of early 2025, the latest information I have):

---

## 1. vLLM

**Short answer: Not practical for custom small models.**

- vLLM is designed for serving large text LLMs (LLaMA, Mistral, GPT-NeoX, etc.). Its core value proposition is **PagedAttention** for efficient KV-cache management when serving many concurrent requests with long sequences.
- It has a **hardcoded model registry**. You cannot simply load an arbitrary `nn.Module`. You must implement a model class that conforms to vLLM's internal interface (`ModelForCausalLM` pattern), register it, and handle vLLM's specific tensor-parallel and KV-cache abstractions.
- The tokenizer/sampling pipeline assumes text tokens — the decoding loop calls a tokenizer to detect stop tokens, EOS, etc. You would need to hack around this for a 1024-vocab codebook.
- **For a 12M param model, vLLM is massive overkill.** The overhead of vLLM's scheduling, PagedAttention, and CUDA graph machinery would likely exceed the cost of running the model itself. vLLM shines at batch sizes of 50+ concurrent requests on 70B+ models.
- **No existing VALL-E or TTS deployments use vLLM** in any public project I'm aware of.

**Verdict: Skip it.** The integration effort is high and the benefit for a small model is zero or negative.

---

## 2. TensorRT / TensorRT-LLM

**TensorRT (general):**
- TensorRT can absolutely optimize a 12M parameter transformer. It doesn't care about vocabulary semantics — a linear projection to 1024 logits is just a matrix multiply.
- You export your PyTorch model to ONNX, then compile with `trtexec` or the Python API. This gives you fused kernels, FP16/INT8 quantization, and optimized CUDA execution.
- **The catch:** You must handle KV-cache manually. TensorRT operates on static graphs, so you need to either (a) export the model with explicit KV-cache inputs/outputs and manage them yourself, or (b) re-run the full sequence each step (slow for autoregressive generation). Option (a) is doable but requires careful ONNX export with dynamic axes.

**TensorRT-LLM:**
- This is NVIDIA's LLM-specific runtime built on TensorRT. It handles KV-cache, in-flight batching, and autoregressive generation natively.
- Like vLLM, it has a **model registry** (LLaMA, GPT, Falcon, etc.). Adding a custom model requires implementing it in TensorRT-LLM's builder API (Python layer that generates TRT engines), which is non-trivial but documented.
- It does not care about token semantics — it just predicts next-token IDs from a vocabulary. A 1024-vocab codebook works fine.
- **However**, TensorRT-LLM is engineered for large models on high-end GPUs. For a 12M model, the framework overhead (scheduling, memory pools, etc.) is disproportionate.

**Verdict:** Plain TensorRT (not TRT-LLM) is a **reasonable option** if you want maximum single-request latency optimization. Export to ONNX, compile to TRT engine, manage KV-cache yourself. For a 12M model, expect sub-millisecond per-step inference.

---

## 3. ONNX Runtime

**Short answer: Best practical option for your use case.**

- ONNX Runtime (ORT) supports transformer decoders with KV-cache via the **`onnxruntime-genai`** package (ONNX Runtime Generative AI). This was introduced specifically for autoregressive generation.
- For a custom model, the workflow is:
  1. Export your PyTorch model to ONNX with explicit `past_key_values` as inputs/outputs
  2. Use `onnxruntime.InferenceSession` to run step-by-step with KV-cache passthrough
  3. Alternatively, `onnxruntime-genai` provides a higher-level `generate()` API, but it currently targets HuggingFace model patterns (GPT-2, LLaMA, Phi)
- **KV-cache support is well-documented.** The ONNX export pattern with `use_cache=True` and dynamic axes for sequence length is a standard pattern.
- ORT supports multiple execution providers: CUDA, TensorRT (as a backend), DirectML, CoreML, OpenVINO, CPU. This gives you deployment flexibility.
- **FP16 and INT8 quantization** are supported via `onnxruntime.quantization`.
- A 12M model will run extremely fast on ORT — likely under 1ms per token on GPU, and potentially viable on CPU too.
- Token semantics are irrelevant to ORT. It runs the computation graph; what the output logits represent is your business.

**Practical pattern for your model:**

```python
# Export with KV-cache
torch.onnx.export(
    model,
    (input_ids, past_key_values),
    "model.onnx",
    input_names=["input_ids", "past_key_0", "past_key_1", ...],
    output_names=["logits", "present_key_0", "present_key_1", ...],
    dynamic_axes={"input_ids": {0: "batch", 1: "seq_len"}, ...}
)

# Run inference
session = ort.InferenceSession("model.onnx", providers=["CUDAExecutionProvider"])
# Loop: feed token + past KV, get logits + new KV
```

**Verdict: Recommended.** Best balance of optimization, flexibility, and low integration effort for a small custom model.

---

## 4. CTranslate2

**Short answer: No meaningful audio model support.**

- CTranslate2 is designed for efficient inference of Transformer models, originally targeting machine translation (Marian, OpenNMT) and later adding Whisper support.
- Supported model types: encoder-decoder (translation), decoder-only (GPT-2, LLaMA via limited support), and Whisper (encoder-decoder for ASR).
- **Whisper support is for speech-to-text**, not text/token-to-speech. It's the only audio-related model.
- There is **no plugin system** for custom architectures. You must use one of the supported model types.
- In principle, if your model architecture matches GPT-2's structure (causal decoder with learned positional embeddings), you could convert it by mapping your weights to CTranslate2's GPT-2 format. But this is fragile and undocumented for non-text use cases.
- The vocabulary/tokenizer layer assumes text. Sampling and decoding utilities are text-oriented.

**Verdict: Not suitable.** Too rigid for custom audio token models.

---

## 5. Other Frameworks Worth Considering

### **llama.cpp / GGML**
- If your model architecture is close to LLaMA (RoPE, RMSNorm, GQA), you could convert it to GGUF format and run it in llama.cpp.
- This gives you optimized CPU inference with SIMD, plus CUDA/Metal offloading.
- The 1024-vocab codebook is fine — llama.cpp doesn't care about token semantics.
- **Downside:** your model needs to match a supported architecture. Custom attention patterns or non-standard components won't work.

### **torch.compile / torch.export (PyTorch native)**
- For a 12M model, **`torch.compile(mode="max-autotune")`** with CUDA graphs may be the simplest high-performance option.
- Combined with `torch.export` for AOT compilation, you get near-TensorRT performance with zero framework integration effort.
- Static KV-cache with `torch.compile` and `torch.export` is supported as of PyTorch 2.2+.
- **This is arguably the lowest-effort, highest-reward option for your specific case.**

### **ExLlamaV2**
- Designed for quantized LLM inference. Not suitable for custom 12M models.

### **Triton Inference Server**
- NVIDIA's model serving framework. It can host ONNX, TensorRT, PyTorch, or TensorFlow models behind a gRPC/HTTP API with dynamic batching.
- This is a **serving layer**, not an optimization layer. Use it in combination with ONNX Runtime or TensorRT as the backend.
- Supports custom Python backends for any model.

---

## 6. How Existing TTS Systems Deploy

### **XTTS (Coqui)**
- Uses a GPT-2-style autoregressive model to predict audio tokens, then a diffusion decoder.
- Deployed via **plain PyTorch inference** with no special framework. The `TTS` Python package just runs `model.forward()` in a loop.
- Some community efforts to export the GPT part to ONNX exist but are unofficial.

### **Parler-TTS (HuggingFace)**
- Uses a decoder-only transformer predicting DAC/EnCodec tokens.
- Deployed via HuggingFace `transformers` with standard `.generate()` — no vLLM, no TensorRT.
- Could be exported to ONNX via `optimum` but no official optimized deployment.

### **VALL-E / VALL-E X**
- Microsoft's research code uses plain PyTorch.
- Open-source reproductions (e.g., `lifeiteng/vall-e`) use PyTorch + `encodec` library.
- **No public deployment uses vLLM or TensorRT-LLM.**

### **Fish Speech / ChatTTS / MeloTTS**
- Fish Speech uses a VQGAN + LLM approach, deployed via plain PyTorch.
- ChatTTS: plain PyTorch.
- MeloTTS: uses ONNX Runtime for some components (specifically the vocoder), but the LM part is plain PyTorch.

### **Sesame / CSM (Conversational Speech Model)**
- Released March 2025. Uses a transformer backbone predicting Mimi codec tokens.
- Deployed via plain PyTorch. The community has created ONNX and MLX ports.

**Key takeaway:** No major TTS system uses vLLM or TensorRT-LLM in production. They all use plain PyTorch or ONNX Runtime.

---

## Recommendation for Your 12M Param Model

**Ranked by practicality:**

1. **`torch.compile` + static KV-cache** — Zero integration work, 2-4x speedup over eager PyTorch, production-ready. Start here.

2. **ONNX Runtime** — If you need cross-platform deployment (e.g., CPU inference, mobile via ONNX Runtime Mobile, or DirectML on Windows). ~1 day of integration work for KV-cache export.

3. **TensorRT** (plain, not TRT-LLM) — If you need absolute minimum latency on NVIDIA GPUs and are willing to manage KV-cache manually. ~2-3 days of integration work.

4. **Triton Inference Server + ONNX/TRT backend** — If you need production serving with dynamic batching and monitoring. Add this on top of option 2 or 3.

Skip vLLM, TensorRT-LLM, CTranslate2, and llama.cpp for this use case.

service_tier: standard

speed: standard

type: assistant