iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Perf][ROCm] Dual-stream decode with hipgraphs (#52033)

Browse Source 
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: Shanshan Shen <467638484@qq.com>
This commit is contained in:
Simon Danielsson
2026-08-30 08:24:20 +00:00
committed by GitHub
co-authored by mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Shanshan Shen
parent 4f78a8fdd0
commit 2c7d7dd64a
2 changed files with 69 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/layers/fused_moe/runner/moe_runner.py
+29 -18
View File
@@ -283,11 +283,24 @@ class MoERunner(MoERunnerInterface):
self._shared_experts: SharedExperts | None = None
if shared_experts is not None:
can_overlap = lambda: self._quant_method.mk_can_overlap_shared_experts
# When unquantized, shared expert inputs alias the hidden states,
# which can lead to race condition in multi-stream mode. Quantized
# routed experts copy the input into a fresh buffer first, breaking
# the alias, so overlap is safe. Only observed on ROCm.
routed_input_is_quantized = lambda: (
self.routed_experts.quant_method.moe_quant_config is not None
and self.routed_experts.quant_method.moe_quant_config.quant_dtype
is not None
)
is_multistream_safe = lambda: (
not current_platform.is_rocm() or routed_input_is_quantized()
)
self._shared_experts = SharedExperts(
shared_experts,
moe_config=moe_config,
enable_dbo=enable_dbo,
mk_can_overlap_shared_experts=can_overlap,
is_multistream_safe=is_multistream_safe,
)
# Needed for string -> MoERunner layer lookup in custom ops.
@@ -582,12 +595,17 @@ class MoERunner(MoERunnerInterface):
router_logits: torch.Tensor,
shared_experts_input: torch.Tensor | None,
input_ids: torch.Tensor | None = None,
shared_experts_overlapping: bool = False,
) -> tuple[torch.Tensor | None, torch.Tensor | UnfinalizedMoEOutput]:
"""Run expert routing and the fused MoE kernel via the quant method.
Orchestrates shared expert execution (before/after), expert selection
via the router, and the actual fused MoE computation. Returns
(shared_expert_output, fused_expert_output).
`shared_experts_overlapping` should be True only if using multi-stream
overlap. Then the shared expert was already launched in a separate
stream, so the results only have to be awaited here.
"""
self._maybe_apply_shared_experts(
shared_experts_input, SharedExpertsOrder.NO_OVERLAP
@@ -617,10 +635,9 @@ class MoERunner(MoERunnerInterface):
shared_experts_input=shared_experts_input,
)
self._maybe_apply_shared_experts(
shared_experts_input,
SharedExpertsOrder.MULTI_STREAM_OVERLAPPED,
)
if shared_experts_overlapping:
assert self._shared_experts is not None
self._shared_experts.wait()
return (
self._shared_experts.output if self._shared_experts is not None else None,
@@ -642,18 +659,6 @@ class MoERunner(MoERunnerInterface):
else nullcontext()
)
def _maybe_sync_shared_experts_stream(
self,
shared_experts_input: torch.Tensor | None,
):
# If router/gate provided, then apply it here.
# (Note: This code runs only when "overlapped mode" is on to allow
# parallel execution of shared experts with the RoutedExperts via
# separate cuda stream)
if self._shared_experts is not None:
assert shared_experts_input is not None
self._shared_experts.maybe_sync_shared_experts_stream(shared_experts_input)
def _maybe_add_zero_expert_output(
self,
result: torch.Tensor,
@@ -878,8 +883,13 @@ class MoERunner(MoERunnerInterface):
# TODO(bnell): this can be removed after MK migration is complete.
self.routed_experts._ensure_moe_quant_config_init()
# Sync aux and main stream for shared expert multi-stream overlap.
self._maybe_sync_shared_experts_stream(shared_experts_input)
# If using multi-stream overlap for shared experts, we must launch it
# before routed expert dispatch.
shared_experts_overlapping = False
if self._shared_experts is not None:
shared_experts_overlapping = self._shared_experts.maybe_forward_async(
shared_experts_input
)
# If the Runner holds the gate, apply it after the stream sync,
# so it can run overlapped with the
@@ -905,6 +915,7 @@ class MoERunner(MoERunnerInterface):
router_logits=router_logits,
shared_experts_input=shared_experts_input,
input_ids=input_ids,
shared_experts_overlapping=shared_experts_overlapping,
)
return self._maybe_combine(
vllm/model_executor/layers/fused_moe/runner/shared_experts.py
+40 -38
View File
@@ -43,6 +43,7 @@ class SharedExperts(torch.nn.Module):
moe_config: FusedMoEConfig,
enable_dbo: bool,
mk_can_overlap_shared_experts: Callable[[], bool],
is_multistream_safe: Callable[[], bool],
):
super().__init__()
@@ -57,6 +58,10 @@ class SharedExperts(torch.nn.Module):
self._mk_can_overlap_shared_experts = mk_can_overlap_shared_experts
# Might not be safe to run multi-stream mode if routed and shared experts
# alias the same inputs
self._is_multistream_safe = is_multistream_safe
# Allow disabling of the separate shared experts stream for
# debug purposes.
# TODO: Remove this after more extensive testings with TP/DP
@@ -65,11 +70,12 @@ class SharedExperts(torch.nn.Module):
logger.debug_once("Disabling MoE shared_experts cuda stream")
self._stream = None
else:
# TODO(rob): enable shared expert overlap with non-cuda-alike.
# aux_stream() returns None on non-cuda-alike platforms.
self._stream = aux_stream()
if self._stream is not None:
logger.debug_once("Enabled separate cuda stream for MoE shared_experts")
# One pair per DBO ubatch id to sync aux and main stream.
self._input_ready_event = [torch.cuda.Event(), torch.cuda.Event()]
self._output_ready_event = [torch.cuda.Event(), torch.cuda.Event()]
# TODO(bnell): Hack for elastic_ep. Get rid of this
def _set_moe_config(self, new_moe_config: FusedMoEConfig):
@@ -106,11 +112,19 @@ class SharedExperts(torch.nn.Module):
if self._mk_can_overlap_shared_experts():
return SharedExpertsOrder.MK_INTERNAL_OVERLAPPED
# On ROCm, empirically only DP-only deployments benefit from the overlap.
overlap_is_beneficial = not current_platform.is_rocm() or (
self._moe_config.moe_parallel_config.dp_size > 1
and self._moe_config.moe_parallel_config.tp_size == 1
)
should_run_shared_in_aux_stream = (
current_platform.is_cuda()
current_platform.is_cuda_alike()
and self._stream is not None
and hidden_states.shape[0]
<= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD
and overlap_is_beneficial
and self._is_multistream_safe()
)
if should_run_shared_in_aux_stream:
@@ -118,38 +132,31 @@ class SharedExperts(torch.nn.Module):
else:
return SharedExpertsOrder.NO_OVERLAP
def maybe_sync_shared_experts_stream(
self,
shared_experts_input: torch.Tensor,
):
experts_order = self._determine_shared_experts_order(shared_experts_input)
def maybe_forward_async(self, shared_experts_input: torch.Tensor) -> bool:
"""Enqueue shared experts on the aux stream without waiting for them.
if experts_order == SharedExpertsOrder.MULTI_STREAM_OVERLAPPED:
assert self._stream is not None
# Record that the clone will be used by shared_experts_stream
# to avoid gc issue from deallocation of hidden_states_clone
# For more details: https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html # noqa: E501
# NOTE: We don't need shared_output.record_stream(current_stream())
# because we synch the streams before using shared_output.
shared_experts_input.record_stream(self._stream)
# Mark sync start point for the aux stream since we will
# run in parallel with router/gate.
self._stream.wait_stream(current_stream())
def _run_in_aux_stream(
self,
shared_experts_input: torch.Tensor,
) -> torch.Tensor:
# TODO: assert that maybe_sync_shared_experts_stream has been called.
# Run shared experts in parallel on a separate stream.
Returns true if the shared experts were enqueued, false otherwise. Call
`wait` to wait for the shared experts to finish if this returns true.
"""
if (
self._determine_shared_experts_order(shared_experts_input)
!= SharedExpertsOrder.MULTI_STREAM_OVERLAPPED
):
return False
assert self._stream is not None
idx = self._output_idx
assert self._output[idx] is None
self._input_ready_event[idx].record(current_stream())
with torch.cuda.stream(self._stream):
output = self._layer(shared_experts_input)
current_stream().wait_stream(self._stream)
self._input_ready_event[idx].wait(self._stream)
self._output[idx] = self._layer(shared_experts_input)
self._output_ready_event[idx].record(self._stream)
return True
return output
def wait(self) -> None:
"""Block the main stream until `maybe_forward_async` output is ready."""
assert self._stream is not None
self._output_ready_event[self._output_idx].wait(current_stream())
@property
def _output_idx(self) -> int:
@@ -174,11 +181,6 @@ class SharedExperts(torch.nn.Module):
assert self._output[self._output_idx] is None
if order == SharedExpertsOrder.MULTI_STREAM_OVERLAPPED:
self._output[self._output_idx] = self._run_in_aux_stream(
shared_experts_input
)
else:
self._output[self._output_idx] = self._layer(shared_experts_input)
self._output[self._output_idx] = self._layer(shared_experts_input)
assert self._output[self._output_idx] is not None