Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion mstar/model/qwen3_tts/submodules.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
4 changes: 4 additions & 0 deletions test/modular/test_qwen3_tts_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading