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 13, 13:15 UTC
[Bugfix] Fix Stream Sync for Shared Expert Overlap (#28430)
Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com> Signed-off-by: Robert Shaw <robertgshaw2@gmail.com> Co-authored-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>
Robert Shaw committed
Nov 11, 2025 at 00:59 UTC
e605e8e3233f895340f46665f93ab37b307491aa
2 files changed
+15
-33
tests/evals/gsm8k/configs/Qwen1.5-MoE-W4A16-CT.yaml
-3
@@ -3,6 +3,3 @@ accuracy_threshold: 0.45
3
num_questions: 1319
4
num_fewshot: 5
5
max_model_len: 4096
6
-# Duo stream incompatabilbe with this model: https://github.com/vllm-project/vllm/issues/28220
7
-env:
8
- VLLM_DISABLE_SHARED_EXPERTS_STREAM: "1"
vllm/model_executor/layers/fused_moe/layer.py
+15
-30
@@ -2456,28 +2456,6 @@ class FusedMoE(CustomOp):
2456
staged_hidden_states.copy_(hidden_states, non_blocking=True)
2457
staged_router_logits.copy_(router_logits, non_blocking=True)
2458
2459
- # If there are shared experts but we are not using a modular kernel,
2460
- # the shared experts must be called here
2461
- if has_separate_shared_experts:
2462
- assert self.shared_experts is not None
2463
-
2464
- if self.shared_experts_stream is not None:
2465
- # For chunked, we start the shared experts stream here
2466
- # (Note that no concurrency with the router/gate)
2467
- self.shared_experts_stream.wait_stream(current_stream())
2468
-
2469
- with torch.cuda.stream(self.shared_experts_stream):
2470
- # Note that staged_hidden_states clone() is necessary
2471
- # here to avoid conflict with the main stream
2472
- shared_output = self.shared_experts(
2473
- staged_hidden_states.clone()
2474
- )
2475
- else:
2476
- shared_output = self.shared_experts(staged_hidden_states)
2477
-
2478
- else:
2479
- shared_output = None
2480
-
2459
# Matrix multiply.
2460
final_hidden_states = self.quant_method.apply(
2461
layer=self,
@@ -2506,11 +2484,7 @@ class FusedMoE(CustomOp):
2484
if has_separate_shared_experts:
2485
assert not isinstance(final_hidden_states, tuple)
2486
assert self.shared_experts is not None
2509
-
2510
- # Here we finish the shared experts stream
2511
- if self.shared_experts_stream is not None:
2512
- current_stream().wait_stream(self.shared_experts_stream)
2513
-
2487
+ shared_output = self.shared_experts(staged_hidden_states)
2488
final_hidden_states = (
2489
shared_output,
2490
final_hidden_states,
@@ -2619,11 +2593,22 @@ class FusedMoE(CustomOp):
2593
assert self.shared_experts is not None
2594
2595
if self.shared_experts_stream is not None:
2596
+ # Clone BEFORE switching streams to avoid race condition
2597
+ # where routed_expert kernel may mutate hidden_states.
2598
+ hidden_states_clone = hidden_states.clone()
2599
+ self.shared_experts_stream.wait_stream(current_stream())
2600
+
2601
# Run shared experts in parallel on a separate stream
2602
with torch.cuda.stream(self.shared_experts_stream):
2624
- # Note that hidden_states clone() is necessary here to avoid
2625
- # conflict with the main stream
2626
- shared_output = self.shared_experts(hidden_states.clone())
2603
+ shared_output = self.shared_experts(hidden_states_clone)
2604
+
2605
+ # Record that the clone will be used by shared_experts_stream
2606
+ # to avoid gc issue from deallocation of hidden_states_clone
2607
+ # For more details: https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html # noqa: E501
2608
+ # NOTE: we dont need shared_output.record_stream(current_stream())
2609
+ # because we synch the streams before using shared_output.
2610
+ hidden_states_clone.record_stream(self.shared_experts_stream)
2611
+
2612
else:
2613
shared_output = self.shared_experts(hidden_states)
2614
else: