Repository navigation
[MUSA] Keep predictor graphs and float32 RMSNorm fallback - #1
yingzhou-bjtu wants to merge 8 commits into
Conversation
723dab4 to
77c65ea
Compare
1fa11c1 to
3be5877
Compare
c88a683 to
fc7b40e
Compare
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.
3fecef3 to
3cb4fcf
Compare
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
changed the title
| tp_size = get_parallel().tp_size | ||
| # note (yingzhou): runtime context is unpublished at construction, | ||
| # so read server_args instead of get_parallel()/get_schedule(). | ||
| tp_size = server_args.tp_size |
There was a problem hiding this comment.
Q: The runtime context is already published here, so can the runner keep get_parallel() and get_schedule() as on main? The raw fields can miss resolved values such as the derived mem_fraction_static and encoder reserve.
| return x | ||
| # note (yingzhou): MUSA float32 has no fused RMSNorm kernel, so keep the | ||
| # CUDA fused path and only native-fallback this dtype on MUSA. | ||
| if x.device.type == "musa" and x.dtype == torch.float32: |
There was a problem hiding this comment.
Q: Does this MUSA float32 RMSNorm fallback run on the pinned SGLang version, or does MUSA dispatch through forward_musa before reaching forward_cuda?
| isolate_sampling=isolate_sampling, | ||
| ) | ||
|
|
||
| def _musa_inference_context(self): |
There was a problem hiding this comment.
Q: Should the repeated MUSA checks become one platform hook beside graph_capture_attention(), with nullcontext by default and inference mode on MUSA?
| # note (yingzhou): MUSA capture_begin rejects inplace updates to | ||
| # inference tensors when capture starts under no_grad. CUDA stays on | ||
| # the original path; MUSA capture and warmup share inference mode. | ||
| capture_mode = ( |
There was a problem hiding this comment.
Q: Which tensor does capture_begin reject here, and does main with sgl-project#2167 fail the same MUSA decode capture without this change?
| ) | ||
|
|
||
| def execute(self, scheduler_output: Any) -> ModelRunnerOutput: | ||
| with self._musa_inference_context(): |
There was a problem hiding this comment.
Q: Can the inference-mode context stay directly inside execute(), and likewise inside execute_launch() and execute_resolve(), instead of using one-use implementation wrappers?
| for _ in range(2) | ||
| ] | ||
| # note (yingzhou): MUSA execute stays in inference mode, but these | ||
| # pinned host buffers are copied into after execute returns. |
There was a problem hiding this comment.
Q: If execute remains in inference mode, can the existing pinned staging allocator be reused instead of making a second copy outside inference mode?
| ) -> Tuple[torch.Tensor, torch.Tensor]: | ||
| # note (yingzhou): MUSA capture stores graph buffers as inference | ||
| # tensors, so replay inplace copies must stay in inference mode. | ||
| replay_mode = ( |
There was a problem hiding this comment.
Q: Replay already runs inside execute's MUSA inference-mode boundary, and the startup buffers are allocated before capture mode. Can the replay-specific inference-mode context be removed?
| ) | ||
| return kwargs | ||
|
|
||
| def forward(self, *args, **kwargs): |
There was a problem hiding this comment.
Q: forward() is reached from execute, which already enters MUSA inference mode, and DLLM does not run on MUSA. Can this override be dropped?
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.
yingzhou-bjtu
changed the title
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
Overview
Enable the MUSA Qwen3-TTS predictor Graph path while keeping the ownership boundary between SGLang-Omni and the native MUSA backend explicit.
Changes
MUSAOmniPlatform.enable_tts_predictor_graph() -> Falseso MUSA inherits the Omni platform default and keeps predictor Graphs enabled.Cross-repository relationship
The native generator-state fix required by MUSA Graph capture is implemented in the companion
torch_musa#1. That PR allocatesseed_extragraph_andoffset_extragraph_outside inference mode so later graph registration can reset them in place.This PR owns the downstream Qwen3-TTS platform behavior;
torch_musa#1owns the native RNG tensor lifetime. The related upstream graph-inference work is sgl-project/sglang-omni#2168.Validation
ad727542.200with a non-emptyaudio/wavresponse of7,864,364bytes.musa graph: Trueduring decode.evidence/qwen3-tts-20260920/ad727542/.