From 181297374ed253d587b0369f1e321e34c39f2184 Mon Sep 17 00:00:00 2001 From: zichen-kuuga Date: Mon, 28 Sep 2026 06:28:14 +0000 Subject: [PATCH 1/3] perf: rotate Model Runner V2 output copy streams --- ...rotate-V2-runner-output-copy-streams.patch | 48 +++++++++++++++++++ vllm_musa/patches/series/README.md | 6 ++- 2 files changed, 52 insertions(+), 2 deletions(-) create mode 100644 vllm_musa/patches/series/0172-MUSA-rotate-V2-runner-output-copy-streams.patch diff --git a/vllm_musa/patches/series/0172-MUSA-rotate-V2-runner-output-copy-streams.patch b/vllm_musa/patches/series/0172-MUSA-rotate-V2-runner-output-copy-streams.patch new file mode 100644 index 000000000000..ed32a02ba6e2 --- /dev/null +++ b/vllm_musa/patches/series/0172-MUSA-rotate-V2-runner-output-copy-streams.patch @@ -0,0 +1,48 @@ +From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001 +From: musa +Date: Mon, 28 Sep 2026 12:00:00 +0800 +Subject: [PATCH] MUSA: rotate V2 runner output copy streams + +On MUSA, once step N+1's copy-stream wait is queued, step N's output copy +event completes only after step N+1 finishes, so each sampled token is +published one step late and the host launches every other step after the +GPU has drained. Give each in-flight step its own output copy stream. + +diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py +--- a/vllm/v1/worker/gpu/model_runner.py ++++ b/vllm/v1/worker/gpu/model_runner.py +@@ -192,6 +192,13 @@ + self.is_encoder_decoder = self.model_config.is_encoder_decoder + + self.output_copy_stream = torch.cuda.Stream(self.device) ++ # MUSA: a copy-stream wait queued for step N+1 delays step N's copy event ++ # until step N+1 finishes, so each in-flight step gets its own stream. ++ self._output_copy_streams = [ ++ self.output_copy_stream, ++ torch.cuda.Stream(self.device), ++ ] ++ self._output_copy_stream_idx = 0 + + # Pipeline parallelism. + self.use_pp = self.parallel_config.pipeline_parallel_size > 1 +@@ -1818,7 +1825,7 @@ + sampler_output=sampler_output, + num_sampled_tokens=num_sampled, + main_stream=self.main_stream, +- copy_stream=self.output_copy_stream, ++ copy_stream=self._next_output_copy_stream(), + check_ep_fault=self.check_ep_fault, + routed_experts=routed_experts, + ) +@@ -1894,6 +1901,11 @@ + + return async_output + ++ def _next_output_copy_stream(self) -> torch.cuda.Stream: ++ streams = self._output_copy_streams ++ self._output_copy_stream_idx = (self._output_copy_stream_idx + 1) % len(streams) ++ return streams[self._output_copy_stream_idx] ++ + def take_draft_token_ids(self) -> DraftTokenIds | None: + return self.draft_tokens_handler.get_draft_tokens() + diff --git a/vllm_musa/patches/series/README.md b/vllm_musa/patches/series/README.md index 08ac75f150d2..4ee544245da0 100644 --- a/vllm_musa/patches/series/README.md +++ b/vllm_musa/patches/series/README.md @@ -17,7 +17,7 @@ is pre-patched. Author headers are normalized to the synthetic `musa ` identity. -Currently **171 patches**. This branch includes the Qwen3.6 patches for common +Currently **172 patches**. This branch includes the Qwen3.6 patches for common GDN decode metadata reuse, uniform-decode SSM slot-mapping removal, and the BF16 W1 tile specialization, plus the contract-bound DeepSeek-V4 MTP sparse-prefill headroom and mixed-prefill queue-fence patches. It additionally @@ -46,7 +46,9 @@ shape fallback. On that same native path the compressor also writes packed kv/score+ape into the state cache so Triton `save_partial_states` is skipped. Interleaved MRoPE rebuilds the T/H/W frequency layout with a strided copy rather than per-channel index arithmetic, which keeps the -per-layer launch count off the eager prefill path. The fused TileLang `hc_head` is +per-layer launch count off the eager prefill path. Model Runner V2 rotates two +output copy streams so a step's sampled tokens are not held back by the next +step's copy-stream wait. The fused TileLang `hc_head` is enabled on MUSA by importing TileLang before the eager JIT decorators capture their module globals. DeepEP shutdown now drops cached handles before native teardown and supports both explicit `destroy()` and legacy destructor-only From fcf25c3bdb10dfc0c40ee70aee63d076d70937b8 Mon Sep 17 00:00:00 2001 From: "zichen.wang" Date: Mon, 28 Sep 2026 15:46:52 +0800 Subject: [PATCH 2/3] perf: size V2 output copy streams by max_concurrent_batches --- ...rotate-V2-runner-output-copy-streams.patch | 43 +++++++++++++------ 1 file changed, 30 insertions(+), 13 deletions(-) diff --git a/vllm_musa/patches/series/0172-MUSA-rotate-V2-runner-output-copy-streams.patch b/vllm_musa/patches/series/0172-MUSA-rotate-V2-runner-output-copy-streams.patch index ed32a02ba6e2..589a0dda21ab 100644 --- a/vllm_musa/patches/series/0172-MUSA-rotate-V2-runner-output-copy-streams.patch +++ b/vllm_musa/patches/series/0172-MUSA-rotate-V2-runner-output-copy-streams.patch @@ -3,29 +3,37 @@ From: musa Date: Mon, 28 Sep 2026 12:00:00 +0800 Subject: [PATCH] MUSA: rotate V2 runner output copy streams -On MUSA, once step N+1's copy-stream wait is queued, step N's output copy -event completes only after step N+1 finishes, so each sampled token is -published one step late and the host launches every other step after the -GPU has drained. Give each in-flight step its own output copy stream. +On MUSA, once a later step's copy-stream wait is queued, an earlier step's +output copy event completes only after that later step finishes, so each +sampled token is published one step late and the host launches every other +step after the GPU has drained. Give each in-flight step its own output copy +stream, sized by max_concurrent_batches, for both the sampling and pooling +outputs. +--- + vllm/v1/worker/gpu/model_runner.py | 18 ++++++++++++++++-- + 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py +index 79953a7..02ae53a 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py -@@ -192,6 +192,13 @@ +@@ -192,6 +192,15 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.is_encoder_decoder = self.model_config.is_encoder_decoder self.output_copy_stream = torch.cuda.Stream(self.device) -+ # MUSA: a copy-stream wait queued for step N+1 delays step N's copy event -+ # until step N+1 finishes, so each in-flight step gets its own stream. -+ self._output_copy_streams = [ -+ self.output_copy_stream, -+ torch.cuda.Stream(self.device), -+ ] ++ # MUSA: a copy-stream wait queued for a later step delays an earlier step's ++ # copy event until that later step finishes, so each in-flight step gets ++ # its own stream. ++ self._output_copy_streams = [self.output_copy_stream] ++ self._output_copy_streams.extend( ++ torch.cuda.Stream(self.device) ++ for _ in range(max(1, vllm_config.max_concurrent_batches - 1)) ++ ) + self._output_copy_stream_idx = 0 # Pipeline parallelism. self.use_pp = self.parallel_config.pipeline_parallel_size > 1 -@@ -1818,7 +1825,7 @@ +@@ -1818,7 +1827,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): sampler_output=sampler_output, num_sampled_tokens=num_sampled, main_stream=self.main_stream, @@ -34,7 +42,7 @@ diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runne check_ep_fault=self.check_ep_fault, routed_experts=routed_experts, ) -@@ -1894,6 +1901,11 @@ +@@ -1894,6 +1903,11 @@ class GPUModelRunner(LoRAModelRunnerMixin): return async_output @@ -46,3 +54,12 @@ diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runne def take_draft_token_ids(self) -> DraftTokenIds | None: return self.draft_tokens_handler.get_draft_tokens() +@@ -1935,7 +1949,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): + pooler_output=pooler_output, + finished_mask=finished_mask, + main_stream=self.main_stream, +- copy_stream=self.output_copy_stream, ++ copy_stream=self._next_output_copy_stream(), + ) + + self.postprocess_num_computed_tokens(input_batch) From 97484d2d8b9e7fdc7a57950e34a3887e39c19c4f Mon Sep 17 00:00:00 2001 From: "zichen.wang" Date: Mon, 28 Sep 2026 16:29:57 +0800 Subject: [PATCH 3/3] perf: use the first output copy stream for the first step --- ...-MUSA-rotate-V2-runner-output-copy-streams.patch | 13 +++++++------ vllm_musa/patches/series/README.md | 6 +++--- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/vllm_musa/patches/series/0172-MUSA-rotate-V2-runner-output-copy-streams.patch b/vllm_musa/patches/series/0172-MUSA-rotate-V2-runner-output-copy-streams.patch index 589a0dda21ab..dd0fb6359081 100644 --- a/vllm_musa/patches/series/0172-MUSA-rotate-V2-runner-output-copy-streams.patch +++ b/vllm_musa/patches/series/0172-MUSA-rotate-V2-runner-output-copy-streams.patch @@ -10,11 +10,11 @@ step after the GPU has drained. Give each in-flight step its own output copy stream, sized by max_concurrent_batches, for both the sampling and pooling outputs. --- - vllm/v1/worker/gpu/model_runner.py | 18 ++++++++++++++++-- - 1 file changed, 16 insertions(+), 2 deletions(-) + vllm/v1/worker/gpu/model_runner.py | 19 +++++++++++++++++-- + 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py -index 79953a7..02ae53a 100644 +index 79953a7..aa11c25 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -192,6 +192,15 @@ class GPUModelRunner(LoRAModelRunnerMixin): @@ -42,19 +42,20 @@ index 79953a7..02ae53a 100644 check_ep_fault=self.check_ep_fault, routed_experts=routed_experts, ) -@@ -1894,6 +1903,11 @@ class GPUModelRunner(LoRAModelRunnerMixin): +@@ -1894,6 +1903,12 @@ class GPUModelRunner(LoRAModelRunnerMixin): return async_output + def _next_output_copy_stream(self) -> torch.cuda.Stream: + streams = self._output_copy_streams ++ stream = streams[self._output_copy_stream_idx] + self._output_copy_stream_idx = (self._output_copy_stream_idx + 1) % len(streams) -+ return streams[self._output_copy_stream_idx] ++ return stream + def take_draft_token_ids(self) -> DraftTokenIds | None: return self.draft_tokens_handler.get_draft_tokens() -@@ -1935,7 +1949,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): +@@ -1935,7 +1950,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): pooler_output=pooler_output, finished_mask=finished_mask, main_stream=self.main_stream, diff --git a/vllm_musa/patches/series/README.md b/vllm_musa/patches/series/README.md index 4ee544245da0..2a1bc46393d7 100644 --- a/vllm_musa/patches/series/README.md +++ b/vllm_musa/patches/series/README.md @@ -46,9 +46,9 @@ shape fallback. On that same native path the compressor also writes packed kv/score+ape into the state cache so Triton `save_partial_states` is skipped. Interleaved MRoPE rebuilds the T/H/W frequency layout with a strided copy rather than per-channel index arithmetic, which keeps the -per-layer launch count off the eager prefill path. Model Runner V2 rotates two -output copy streams so a step's sampled tokens are not held back by the next -step's copy-stream wait. The fused TileLang `hc_head` is +per-layer launch count off the eager prefill path. Model Runner V2 gives each +in-flight step its own output copy stream so a step's sampled tokens are not +held back by a later step's copy-stream wait. The fused TileLang `hc_head` is enabled on MUSA by importing TileLang before the eager JIT decorators capture their module globals. DeepEP shutdown now drops cached handles before native teardown and supports both explicit `destroy()` and legacy destructor-only