From ecd833cb89acf8ef6afce985547d058b933e88cb Mon Sep 17 00:00:00 2001 From: mARTin-B78 Date: Sun, 11 Oct 2026 19:18:31 +0200 Subject: [PATCH] feat: opt-in torch.compile of the code predictor loop (TTS_COMPILE_PREDICTOR=1) Co-Authored-By: Claude Sonnet 5.5 --- Dockerfile | 4 ++-- patches/predictor_compile.patch | 25 +++++++++++++++++++++++++ 2 files changed, 27 insertions(+), 2 deletions(-) create mode 100644 patches/predictor_compile.patch diff --git a/Dockerfile b/Dockerfile index 9b37f14..6327555 100644 --- a/Dockerfile +++ b/Dockerfile @@ -24,8 +24,8 @@ RUN git clone https://github.com/andimarafioti/faster-qwen3-tts.git /app && \ cd /app && git checkout --detach ${FASTER_QWEN3_TTS_REF} # Apply DGX Spark patches (non_streaming_mode=True, per-voice temperature/top_k/top_p) -COPY patches/openai_server.patch /tmp/ -RUN cd /app && git apply /tmp/openai_server.patch +COPY patches/openai_server.patch patches/predictor_compile.patch /tmp/ +RUN cd /app && git apply /tmp/openai_server.patch && git apply /tmp/predictor_compile.patch # Create virtual environment (Ubuntu 24.04 enforces PEP 668) ENV VIRTUAL_ENV=/opt/venv diff --git a/patches/predictor_compile.patch b/patches/predictor_compile.patch new file mode 100644 index 0000000..71ee18b --- /dev/null +++ b/patches/predictor_compile.patch @@ -0,0 +1,25 @@ +diff --git a/faster_qwen3_tts/predictor_graph.py b/faster_qwen3_tts/predictor_graph.py +index 38a52ca..70851eb 100644 +--- a/faster_qwen3_tts/predictor_graph.py ++++ b/faster_qwen3_tts/predictor_graph.py +@@ -13,6 +13,8 @@ Strategy: + - Unroll the full 15-step loop for deterministic shapes + - Capture the entire loop as a single CUDA graph + """ ++import os ++ + import torch + from transformers import StaticCache + from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask +@@ -175,6 +177,11 @@ class PredictorGraph: + self._init_cache_layers() + self._build_attention_masks() + ++ # Opt-in: fuse the 15-step predictor loop with Inductor before graph capture. ++ # ~17% faster streaming on GB10 (0.6B); adds ~1 min to the first warmup. ++ if os.environ.get("TTS_COMPILE_PREDICTOR") == "1": ++ self._full_loop = torch.compile(self._full_loop, dynamic=False) ++ + for _ in range(num_warmup): + self.static_cache.reset() + self._full_loop()