From 91dff7d4e58c32457522a12333dc143241b4bf02 Mon Sep 17 00:00:00 2001 From: merceod <57202593+merceod@users.noreply.github.com> Date: Mon, 21 Sep 2026 12:48:54 +0000 Subject: [PATCH] Capture the eager kernels for the qwen3tts talker decode --- mstar/model/qwen3_tts/submodules.py | 7 ++++++- test/modular/test_qwen3_tts_model.py | 4 ++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/mstar/model/qwen3_tts/submodules.py b/mstar/model/qwen3_tts/submodules.py index ad3b99bbc..2dde9468e 100644 --- a/mstar/model/qwen3_tts/submodules.py +++ b/mstar/model/qwen3_tts/submodules.py @@ -708,7 +708,12 @@ def get_cuda_graph_configs( }, ), capture_batch_sizes=self.DECODE_CAPTURE_BATCH_SIZES, - compile=True, + # Capture the eager kernels. A torch.compile'd capture picks its + # kernels by benchmarking in each process, and sixty sampling + # steps turn the resulting rounding differences into a different + # token sequence per server instance for the same seed. The eager + # kernels are the same everywhere, and the capture is faster. + compile=False, )] def get_piecewise_cuda_graph_configs( diff --git a/test/modular/test_qwen3_tts_model.py b/test/modular/test_qwen3_tts_model.py index eed6b957d..dfb87ead5 100644 --- a/test/modular/test_qwen3_tts_model.py +++ b/test/modular/test_qwen3_tts_model.py @@ -702,6 +702,10 @@ def test_qwen3_tts_talker_batches_and_captures_decode(): graph_config = submodule.get_cuda_graph_configs(torch.device("cpu"))[0] assert graph_config.capture_graph_walk == "talker_decode" assert graph_config.capture_batch_sizes == [1, 2, 4, 8, 16, 32] + # The decode graph captures the eager kernels: a compiled capture picks + # kernels by benchmarking per process, and the same seed then gives a + # different token sequence on every server instance. + assert graph_config.compile is False assert graph_config.single_request_inputs.tensor_inputs[ "suppress_eos" ].item() is True