Repository navigation
[MUSA] Enable TTS predictor graphs - #2168
Merged
Ratish1 merged 8 commits intoSep 21, 2026
Merged
Conversation
This was referenced Sep 14, 2026
20 of 29 tasks
yeahdongcn
requested changes
Sep 15, 2026
Contributor
Author
|
/tag-and-rerun-ci |
yingzhou-bjtu
force-pushed
the
upstream-pr1-musa-common
branch
from
September 17, 2026 06:39
254505f to
fc7b40e
Compare
added 5 commits
September 18, 2026 14:39
Keep current main's predictor graph backend. Wrap MUSA capture/replay in inference mode, read runner config from server_args before runtime context is published, and stop disabling MUSA predictor graphs.
Keep the current-main graph backend rebase, and apply the existing fork PR1 comment/import cleanup so capture comments sit next to the MUSA inference-mode guard.
MUSA graph replay leaves logits as inference tensors. Wrap execute, launch, and resolve in inference mode so codec suppress can write those logits in place, and allocate pinned host staging outside inference mode so later host copies stay valid.
Keep the MUSA execute path in inference mode and assert host staging stays an ordinary tensor after execute returns.
Restore the capture decorator so CUDA stays on no_grad while MUSA still enters inference mode through capture_mode. Collapse duplicated host staging allocation and sign the constructor comment.
yingzhou-bjtu
force-pushed
the
upstream-pr1-musa-common
branch
from
September 18, 2026 10:25
fc7b40e to
18d9b12
Compare
Ratish1
reviewed
Sep 18, 2026
added 3 commits
September 20, 2026 15:27
Keep Omni-owned MUSA pieces: predictor graph enablement, constructor server_args reads, and the float32 RMSNorm native fallback. Graph RNG capture_begin now belongs in torch_musa, so Omni no longer wraps capture, replay, or execute in inference mode.
ModelWorker publishes server_args before constructing SGLModelRunner, so keep get_parallel() and get_schedule() as on main. Raw server_args can miss the derived mem_fraction_static after encoder reserve.
yingzhou-bjtu
changed the title
Ratish1
approved these changes
Sep 21, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
MUSA currently disables Qwen3-TTS predictor Graphs by overriding
enable_tts_predictor_graph()toFalse. CUDA already inherits theOmniPlatform default (
True). This PR deletes that override.Modifications
MUSAOmniPlatform.enable_tts_predictor_graph() -> False.tests/unit_test/test_platforms.py.This PR is based on
main.Related Issues
Native Graph RNG fix;
torch_musa 2.11.0.post2already supports this:yingzhou-bjtu/torch_musa#1
Qwen3-TTS MUSA device-path adaptation:
#2167
Accuracy Test
N/A. The change is a platform default.
Benchmark & Profiling
N/A.
Validation
Static checks at head
ad727542:Both passed.
Checklist
CI
CI runs on self-hosted GPU runners and requires a maintainer to add the
run-cilabel. Once labeled, every subsequent push re-triggers CI aslong as the label remains. Use
/tag-and-rerun-ci higgsor/tag-and-rerun-ci mossor/tag-and-rerun-ci qwen3-ttsto select a TTS CImodel, and
/tag-and-rerun-ci fun-asr,/tag-and-rerun-ci qwen3-asror/tag-and-rerun-ci whisper-asrto select an ASR CI model. One selector fromeach family can be combined, for example
/tag-and-rerun-ci moss fun-asr.Draft PRs are skipped even if labeled.