A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 01:15 UTC
[Model Bash][DeepSeekR1] Remove Shared Expert Clone (#34344)
Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com>
Robert Shaw committed
Feb 19, 2026 at 10:56 UTC
4685a630a293cd7c928092efd0f8c2606a770877
2 files changed
+11
-13
vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py
+10
-12
@@ -240,24 +240,22 @@ class DefaultMoERunner(MoERunner):
240
)
241
)
242
243
- hidden_states_clone: torch.Tensor | None = None
243
+ shared_experts_input: torch.Tensor | None = None
244
if use_shared_experts_stream:
245
assert self.shared_experts_stream is not None
246
+ assert self.moe_config.disable_inplace
247
248
shared_experts_input = (
249
shared_input if shared_input is not None else hidden_states
250
)
251
251
- # Clone BEFORE switching streams to avoid race condition
252
- # where routed_expert kernel may mutate hidden_states.
253
- hidden_states_clone = shared_experts_input.clone()
254
-
255
- # Record that the clone will be used by shared_experts_stream
256
- # to avoid gc issue from deallocation of hidden_states_clone
257
- # For more details: https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html # noqa: E501
252
+ # Record that the shared_experts_input will be used in the
253
+ # shared_experts_stream to to avoid gc issue from
254
+ # deallocation. For more details:
255
+ # https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html # noqa: E501
256
# NOTE: We don't need shared_output.record_stream(current_stream())
257
# because we synch the streams before using shared_output.
260
- hidden_states_clone.record_stream(self.shared_experts_stream)
258
+ shared_experts_input.record_stream(self.shared_experts_stream)
259
260
# Mark sync start point for the separate shared experts
261
# stream here since we want to run in parallel with the
@@ -265,7 +263,7 @@ class DefaultMoERunner(MoERunner):
263
assert self.shared_experts_stream is not None
264
self.shared_experts_stream.wait_stream(current_stream())
265
268
- return use_shared_experts_stream, hidden_states_clone
266
+ return use_shared_experts_stream, shared_experts_input
267
268
def ensure_dp_chunking_init(self):
269
if not self.use_dp_chunking or self.batched_hidden_states is not None:
@@ -584,7 +582,7 @@ class DefaultMoERunner(MoERunner):
582
583
use_chunked_impl = self.use_dp_chunking
584
587
- use_shared_experts_stream, hidden_states_clone = (
585
+ use_shared_experts_stream, shared_experts_input = (
586
self._maybe_setup_shared_experts_stream(
587
hidden_states,
588
shared_input,
@@ -726,7 +724,7 @@ class DefaultMoERunner(MoERunner):
724
with torch.cuda.stream(self.shared_experts_stream):
725
# Note that hidden_states clone() is necessary here to avoid
726
# conflict with the main stream
729
- shared_output = self.shared_experts(hidden_states_clone)
727
+ shared_output = self.shared_experts(shared_experts_input)
728
current_stream().wait_stream(self.shared_experts_stream)
729
730
final_hidden_states = (
vllm/model_executor/models/minicpm.py
+1
-1
@@ -175,7 +175,7 @@ class MiniCPMMoE(nn.Module):
175
)
176
177
final_hidden_states = fused_experts(
178
- hidden_states, self.ws, self.w2s, topk_weights, topk_ids, inplace=True
178
+ hidden_states, self.ws, self.w2s, topk_weights, topk_ids, inplace=False
179
)
180
181
if self.tp_size > 1: