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 12, 22:15 UTC
[Bugfix] Fix precision corruption when shared_experts_stream=None (#28942)
Signed-off-by: zhyajie <yajizhan@amd.com> Co-authored-by: zhyajie <yajizhan@amd.com>
杰兮 committed
Nov 20, 2025 at 03:30 UTC
9d2d5612573c20f8bf00242a8525c2a5dcfe4c06
2 files changed
+8
-6
vllm/model_executor/layers/fused_moe/layer.py
+7
-4
@@ -371,8 +371,8 @@ class FusedMoE(CustomOp):
371
logger.info_once("Disabling MoE shared_experts cuda stream")
372
self.shared_experts_stream = None
373
else:
374
- # TODO(rob): enable shared expert overlap with non-cuda.
375
- # aux_stream() returns None on non-cuda platforms.
374
+ # TODO(rob): enable shared expert overlap with non-cuda-alike.
375
+ # aux_stream() returns None on non-cuda-alike platforms.
376
self.shared_experts_stream = aux_stream()
377
if self.shared_experts_stream is not None:
378
logger.info_once("Enabled separate cuda stream for MoE shared_experts")
@@ -1865,6 +1865,11 @@ class FusedMoE(CustomOp):
1865
hidden_states_combined, router_logits = get_ep_group().dispatch(
1866
hidden_states, router_logits, self.is_sequence_parallel
1867
)
1868
+ # Run shared experts before matrix multiply.
1869
+ # because matrix multiply maybe modify the hidden_states.
1870
+ if has_separate_shared_experts and not use_shared_experts_stream:
1871
+ assert self.shared_experts is not None
1872
+ shared_output = self.shared_experts(hidden_states)
1873
1874
# Matrix multiply.
1875
final_hidden_states = self.quant_method.apply(
@@ -1908,8 +1913,6 @@ class FusedMoE(CustomOp):
1913
# conflict with the main stream
1914
shared_output = self.shared_experts(hidden_states_clone)
1915
current_stream().wait_stream(self.shared_experts_stream)
1911
- else:
1912
- shared_output = self.shared_experts(hidden_states)
1916
1917
final_hidden_states = (
1918
shared_output,
vllm/utils/torch_utils.py
+1
-2
@@ -426,8 +426,7 @@ def aux_stream() -> torch.cuda.Stream | None:
426
427
from vllm.platforms import current_platform
428
429
- # TODO: validate this works properly on ROCm platform.
430
- if _aux_stream is None and current_platform.is_cuda():
429
+ if _aux_stream is None and current_platform.is_cuda_alike():
430
_aux_stream = torch.cuda.Stream()
431
432
return _aux_stream