26 lines
1.0 KiB
Diff
26 lines
1.0 KiB
Diff
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()
|